StableDiffusion example doesn't currently work without adjustment

[mdvthu]

Issue Type

Documentation Bug

Source

binary

Keras Version

keras_cv 0.9.0

Custom Code

No

OS Platform and Distribution

macOS, Windows, Ubuntu

Python version

3.12.4

GPU model and memory

CPU, Nvidia 3080, and M1 (METAL)

Current Behavior?

Following the example verbatim from Keras documentation pages including https://keras.io/guides/keras_cv/generate_images_with_stable_diffusion/ produces shape mismatch errors:

ValueError: Exception encountered when calling DiffusionModelV2.call().

Invalid input shape for input Tensor("data_2:0", shape=(1, 77, 1024), dtype=float32). Expected shape (None, 96, 96, 4), but input has incompatible shape (1, 77, 1024)

Arguments received by DiffusionModelV2.call():
  • inputs={'latent': 'tf.Tensor(shape=(1, 96, 96, 4), dtype=float32)', 'timestep_embedding': 'tf.Tensor(shape=(1, 320), dtype=float32)', 'context': 'tf.Tensor(shape=(1, 77, 1024), dtype=float32)'}
  • training=False
  • mask={'latent': 'None', 'timestep_embedding': 'None', 'context': 'None'}

Standalone code to reproduce the issue or tutorial link

Follow the documentation on https://keras.io/guides/keras_cv/generate_images_with_stable_diffusion/


python3 -m venv venv
. ./venv/bin/activate
python3 -m pip install tensorflow keras_cv IPython
python3 -m IPython


```python
import keras_cv
model = keras_cv.models.StableDiffusion(
    img_width=512, img_height=512, jit_compile=False
)
images = model.text_to_image("photograph of an astronaut riding a horse", batch_size=3)

### Relevant log output

```shell
ValueError: Exception encountered when calling DiffusionModelV2.call().

Invalid input shape for input Tensor("data_2:0", shape=(1, 77, 1024), dtype=float32). Expected shape (None, 96, 96, 4), but input has incompatible shape (1, 77, 1024)

Arguments received by DiffusionModelV2.call():
  • inputs={'latent': 'tf.Tensor(shape=(1, 96, 96, 4), dtype=float32)', 'timestep_embedding': 'tf.Tensor(shape=(1, 320), dtype=float32)', 'context': 'tf.Tensor(shape=(1, 77, 1024), dtype=float32)'}
  • training=False
  • mask={'latent': 'None', 'timestep_embedding': 'None', 'context': 'None'}

[Xisiqomelir]

I noticed the same when running on CoLab. I believe the KerasCV documentation here helps: https://keras.io/keras_cv/

There are currently two ways to install Keras 3 with KerasCV. To install the stable versions of KerasCV and Keras 3, you should install Keras 3 after installing KerasCV. This is a temporary step while TensorFlow is pinned to Keras 2, and will no longer be necessary after TensorFlow 2.16.

So on CoLab which is currently on TensorFlow 2.15 you have to comment out the Keras 3 upgrade and only run

pip install --upgrade keras-cv tensorflow

This then runs until the mixed precision section which gives out of resource errors on the free tier GPUs.

[Eom-taeseon]

I have a same issue in Colab. I upgrade my keras-cv and tensorflow, and after that upgrade keras too. This is my installation code

!pip install keras-cv==0.6.0 -q
!pip install tensorflow==2.4 -q
!pip install keras-core -q

!pip install --upgrade keras-cv tensorflow
!pip install --upgrade keras-cv keras

But I got same error.

ValueError                                Traceback (most recent call last)
[<ipython-input-53-0a64cfdf4765>](https://localhost:8080/#) in <cell line: 1>()
----> 1 generated_image = manhwa_model.text_to_image(
      2     "Yoda", batch_size = 3, unconditional_guidance_scale=40
      3 )

8 frames
[/usr/local/lib/python3.10/dist-packages/keras/src/models/functional.py](https://localhost:8080/#) in _adjust_input_rank(self, flat_inputs)
    242                     adjusted.append(ops.expand_dims(x, axis=-1))
    243                     continue
--> 244             raise ValueError(
    245                 f"Invalid input shape for input {x}. Expected shape "
    246                 f"{ref_shape}, but input has incompatible shape {x.shape}"

ValueError: Exception encountered when calling DiffusionModel.call().

Invalid input shape for input Tensor("data_2:0", shape=(3, 77, 768), dtype=float16). Expected shape (None, 64, 64, 4), but input has incompatible shape (3, 77, 768)

Arguments received by DiffusionModel.call():
  • inputs={'latent': 'tf.Tensor(shape=(3, 64, 64, 4), dtype=float32)', 'timestep_embedding': 'tf.Tensor(shape=(3, 320), dtype=float32)', 'context': 'tf.Tensor(shape=(3, 77, 768), dtype=float16)'}
  • training=False
  • mask={'latent': 'None', 'timestep_embedding': 'None', 'context': 'None'}

How can I solve this?
The URL I referenced : https://keras.io/examples/generative/finetune_stable_diffusion/