Update app.py

app.py

@@ -53,7 +53,7 @@ torch.cuda.empty_cache()
 
 # Inference function
 @spaces.GPU(duration=25)
-def generate_image(
+async def generate_image(
     prompt,
     seed=24,
     width=DEFAULT_WIDTH,
@@ -95,7 +95,6 @@ def generate_image(
                 encoder_hidden_states=static_prompt_embeds,
                 txt_ids=static_text_ids,
                 img_ids=static_latent_image_ids,
-                joint_attention_kwargs=pipe.joint_attention_kwargs,
                 return_dict=False,
             )
     torch.cuda.current_stream().wait_stream(s)
@@ -111,7 +110,6 @@ def generate_image(
             encoder_hidden_states=static_prompt_embeds,
             txt_ids=static_text_ids,
             img_ids=static_latent_image_ids,
-            joint_attention_kwargs=pipe.joint_attention_kwargs,
             return_dict=False,
         )[0]
         static_latents_out = pipe.scheduler.step(
@@ -122,7 +120,7 @@ def generate_image(
         )
 
     # Graph-based generation function
-    def generate_with_graph(
+    async def generate_with_graph(
         latents,
         prompt_embeds,
         pooled_prompt_embeds,
@@ -140,7 +138,7 @@ def generate_image(
         return static_output
 
     # Only generate the last image in the sequence
-    img = pipe.generate_images(
+    img = await pipe.generate_images(
         prompt=prompt,
         width=width,
         height=height,
@@ -251,7 +249,7 @@ with gr.Blocks() as demo:
         outputs=[result, seed, latency],
         show_progress="full",
         api_name="RealtimeFlux",
-        queue=False,
+        queue=False
     )
 
     def update_ui(realtime_enabled):
@@ -269,10 +267,9 @@ with gr.Blocks() as demo:
     )
 
     async def realtime_generation(*args):
+        print("realtime_generation")
         if args[0]:  # If realtime is enabled
-            loop = asyncio.get_event_loop()
-            result = await loop.run_in_executor(None, next, generate_image(*args[1:]))
-            return result
+            return await generate_image(*args[1:])
 
     prompt.submit(
         fn=generate_image,