Update app.py

app.py

@@ -5,10 +5,11 @@ torch.backends.cudnn.allow_tf32 = True
 import gradio as gr
 import numpy as np
 import random
+import spaces
 import time
 from diffusers import DiffusionPipeline, AutoencoderTiny
+from diffusers.models.attention_processor import AttnProcessor2_0
 from custom_pipeline import FluxWithCFGPipeline
-import asyncio
 
 # Constants
 MAX_SEED = np.iinfo(np.int32).max
@@ -49,10 +50,9 @@ if hasattr(pipe, "transformer") and torch.cuda.is_available():
 
 torch.cuda.empty_cache()
 
-
-
-# Inference function (async)
-async def generate_image(
+# Inference function
+@spaces.GPU(duration=25)
+def generate_image(
     prompt,
     seed=24,
     width=DEFAULT_WIDTH,
@@ -118,7 +118,7 @@ async def generate_image(
             static_latents_out, height, width, "pil"
         )
 
-    # Graph-based generation function (synchronous)
+    # Graph-based generation function
     def generate_with_graph(
         latents,
         prompt_embeds,
@@ -136,6 +136,7 @@ async def generate_image(
         g.replay()
         return static_output
 
+    # Only generate the last image in the sequence
     img = pipe.generate_images(
         prompt=prompt,
         width=width,
@@ -264,10 +265,10 @@ with gr.Blocks() as demo:
         concurrency_limit=None,
     )
 
-    async def realtime_generation(*args):
-        print("realtime_generation")
+    def realtime_generation(*args):
         if args[0]:  # If realtime is enabled
-            return await generate_image(*args[1:])
+            img, seed, latency = generate_image(*args[1:])
+            return img, seed, latency
 
     prompt.submit(
         fn=generate_image,