StableDiffusion.text_to_image() casuses an excaption in Colab

[KouichiMatsuda]

Hi Keras Team,

Current Behavior:

The code based on https://keras.io/api/keras_cv/models/tasks/stable_diffusion/ causes an exception: ValueError: Exception encountered when calling DiffusionModelV2.call().

https://colab.research.google.com/drive/1OYet7JBOwgt7L5itxOOzVg-jclgPpdnT?usp=sharing

StableDiffusion class, too.

Am I missing something?

Steps To Reproduce:

https://colab.research.google.com/drive/1OYet7JBOwgt7L5itxOOzVg-jclgPpdnT?usp=sharing

Version:

Keras 3.4.1
TF 2.16.2
KerasCV 0.9.0

[pjogi-testy]

I have the same error when trying to replicate any Stable-Diffusion-related tutorial from official Keras repo (eg. https://keras.io/examples/generative/random_walks_with_stable_diffusion/), no matter if run locally or using original repo in colab. Possibly something is broken with Keras3 and latest TF? It seems like encoder (77 tokens, 768 values) does not communicate with diffusor UNET (basic 64x64x3 shape). Any clues would be most welcome. Or even confirmation on which exact versions of tf, keras, keras_cv does the official repo works on, because the Keras api is so inconsistent between versions that it is really hard to follow on. Further details of Error:

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

Invalid input shape for input Tensor("data_2:0", shape=(3, 77, 768), dtype=float32). 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=float32)'}
• training=False
• mask={'latent': 'None', 'timestep_embedding': 'None', 'context': 'None'}

[heydaari]

i have the same as
#2467 (comment)

i tried different backends , issue wont be gone

[OttoERM]

I got the same problem by following this tensorflow tutorial with the same error as @pjogi-testy

Long story short got it working on: keras 2.13.1, keras-core 0.1.7, keras-cv 0.9.0, tensorflow 2.13.1
And also on: keras 2.15.0, keras-core 0.1.7, keras-cv 0.9.0, tensorflow 2.15.1

import time
import keras_cv
from tensorflow import keras
import matplotlib.pyplot as plt
from PIL import Image

model = keras_cv.models.StableDiffusion(img_width=512, img_height=512)

image = model.text_to_image(prompt="Flower", batch_size=1, num_steps=15)

Image.fromarray(image[0]).save("Flower.png")
print("Saved at flower.png")

Didn't tried out tensorflow 2.16.*
Anyways I guess is just a broken version between keras 3.4.1 and tensorflow 2.17.0 (Latest release at this time)

In the keras repo readme there is a note "Keras 3 will not function with TensorFlow 2.14 or earlier." Not sure how are you suppose to use Keras 3 because whenever I installed tensorflow a specific version of keras was added, I install --upgrade the keras version but gave me a version error incompatibility.