Repository: tensorflow
Issue: [ tf-opt ] Keras Official Implementation of Stable Diffusion fails to generate stablehlo mlir
Labels: stat:awaiting tensorflower,type:bug,comp:keras,TF2.14

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

v2.14.0-rc1-21-g4dacf3f368e 2.14.0

### Custom code

Yes

### OS platform and distribution

Ubuntu 22.04

### Mobile device

_No response_

### Python version

3.10.12

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

After saving the diffusion model in saved_model format using the code given below, I ran the following command to convert the saved_model into an input mlir. 

```
tf-mlir-translate --savedmodel-objectgraph-to-mlir --tf-savedmodel-exported-names=serving_fn ./signed-model/ -o tf_executor.mlir
```

This actually produces a large mlir file of 6.5 GB . Then I ran this: 
``` tf-opt --tf-executor-graph-pruning --tf-executor-to-functional-conversion --canonicalize --tf-lower-to-mlprogram-and-hlo ./tf_executor.mlir -o stablehlo.mlir```

This process runs for a couple of minutes and ends without any error. However, in the ```stablehlo.mlir``` does not have any mlir in it. Its an empty file with the following in it: "module{}". It does not provide any error message at all. 

I was expecting the second command to generate a stablehlo mlir file corresponding to the diffusion model which I could later pass on to the iree-importer for importing into iree. 


### Standalone code to reproduce the issue

```shell
from keras_cv import models
from keras_cv.src.models.stable_diffusion.constants import _ALPHAS_CUMPROD
import tensorflow as tf
from keras_cv.src.backend import ops
import math
from keras_cv.models.stable_diffusion import DiffusionModel

IMG_HEIGHT = 512
IMG_WIDTH = 512
MAX_PROMPT_LENGTH = 77
ALPHAS_CUMPROD_TF = tf.constant(_ALPHAS_CUMPROD)
UNCONDITIONAL_GUIDANCE_SCALE = 7.5
HIDDEN_DIM = 768
SEED = None


signature_dict = {
    "context": tf.TensorSpec(shape=[None, MAX_PROMPT_LENGTH, HIDDEN_DIM], dtype=tf.float32, name="context"),
    "unconditional_context": tf.TensorSpec(
        shape=[None, MAX_PROMPT_LENGTH, HIDDEN_DIM], dtype=tf.float32, name="unconditional_context"
    ),
    "num_steps": tf.TensorSpec(shape=[], dtype=tf.int32, name="num_steps"),
    "batch_size": tf.TensorSpec(shape=[], dtype=tf.int32, name="batch_size"),
}

def diffusion_model_exporter(model: tf.keras.Model):
    @tf.function
    def get_timestep_embedding(timestep, batch_size, dim=320, max_period=10000):
        half = dim // 2
        range = ops.cast(ops.arange(0, half), "float32")
        freqs = ops.exp(-math.log(max_period) * range / half)
        args = ops.convert_to_tensor([timestep], dtype="float32") * freqs
        embedding = ops.concatenate([ops.cos(args), ops.sin(args)], 0)
        embedding = ops.reshape(embedding, [1, -1])
        return ops.repeat(embedding, batch_size, axis=0)
    
    @tf.function(input_signature=[signature_dict])
    def serving_fn(inputs):
        img_height = tf.cast(tf.math.round(IMG_HEIGHT / 128) * 128, tf.int32)
        img_width = tf.cast(tf.math.round(IMG_WIDTH / 128) * 128, tf.int32)

        batch_size = inputs["batch_size"]
        num_steps = inputs["num_steps"]

        context = inputs["context"]
        unconditional_context = inputs["unconditional_context"]

        latent = tf.random.normal((batch_size, img_height // 8, img_width // 8, 4))

        timesteps = tf.range(1, 1000, 1000 // num_steps)
        alphas = tf.map_fn(lambda t: ALPHAS_CUMPROD_TF[t], timesteps, dtype=tf.float32)
        alphas_prev = tf.concat([[1.0], alphas[:-1]], 0)

        index = num_steps - 1
        latent_prev = None
        for timestep in timesteps[::-1]:
            latent_prev = latent
            t_emb = get_timestep_embedding(timestep, batch_size)
            unconditional_latent = model(
                [latent, t_emb, unconditional_context], training=False
            )
            latent = model([latent, t_emb, context], training=False)
            latent = unconditional_latent + UNCONDITIONAL_GUIDANCE_SCALE * (
                latent - unconditional_latent
            )
            a_t, a_prev = alphas[index], alphas_prev[index]
            pred_x0 = (latent_prev - tf.math.sqrt(1 - a_t) * latent) / tf.math.sqrt(a_t)
            latent = (
                latent * tf.math.sqrt(1.0 - a_prev) + tf.math.sqrt(a_prev) * pred_x0
            )
            index = index - 1

        return {"latent": latent}

    return serving_fn

diffuser = DiffusionModel(IMG_HEIGHT, IMG_WIDTH, MAX_PROMPT_LENGTH)

tf.saved_model.save(diffuser, './stable-diffusion/signed-model', signatures={"serving_default": diffusion_model_exporter(diffuser)} )
```


### Relevant log output

_No response_