Update app.py

app.py

@@ -2,7 +2,7 @@ import torch
 import spaces
 import gradio as gr
 from diffusers import DiffusionPipeline
-from optimization import optimize_pipeline_ 
+
 import spaces
 from dataclasses import dataclass
 import json
@@ -25,21 +25,7 @@ pipe = DiffusionPipeline.from_pretrained(
 )
 
 
-#pipe.transformer.set_attention_backend(attention_backend)
-
-# if enable_compile:
-#     print("Compiling transformer...")
-#     pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune-no-cudagraphs", fullgraph=False)
-
-pipe.to("cuda")
-
-optimize_pipeline_(
-    pipe,
-    prompt="prompt",
-    num_inference_steps=1,
-    guidance_scale=0.0
-)
-
+pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
 
 @spaces.GPU
 def generate_image(prompt, height, width, num_inference_steps, seed, randomize_seed):