Update app.py

app.py

@@ -23,8 +23,7 @@ torch.cuda.empty_cache()
 
 # Inference function
 @spaces.GPU(duration=25)
-def generate_image(prompt, seed=42, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, randomize_seed=False, num_inference_steps=2):
-    
+def generate_image(prompt, seed=42, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT, randomize_seed=False, num_inference_steps=2, progress=gr.Progress(track_tqdm=True)):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator().manual_seed(seed)
@@ -43,7 +42,6 @@ def generate_image(prompt, seed=42, width=DEFAULT_WIDTH, height=DEFAULT_HEIGHT,
         latency = f"Latency: {(time.time()-start_time):.2f} seconds"    
         yield img, seed, latency
 
-
 # Example prompts
 examples = [
     "a tiny astronaut hatching from an egg on the moon",
@@ -72,11 +70,13 @@ with gr.Blocks() as demo:
                     show_label=False,
                     container=False,
                 )
+                generateBtn = gr.Button("🖼️ Generate Image", visible=False)
                 enhanceBtn = gr.Button("🚀 Enhance Image")
 
                 with gr.Column("Advanced Options"):
                     with gr.Row():
-                        latency = gr.Text(show_label=False)
+                        realtime = gr.Checkbox(label="Realtime Toggler", info="If TRUE then uses more GPU but create image in realtime.", value=True)
+                        latency = gr.Text(label="Latency")
                     with gr.Row():
                         seed = gr.Number(label="Seed", value=42, precision=0)
                         randomize_seed = gr.Checkbox(label="Randomize Seed", value=False)
@@ -96,26 +96,60 @@ with gr.Blocks() as demo:
                 cache_examples="lazy" 
             )
 
-    # Event handling - Trigger image generation on button click or input change
     enhanceBtn.click(
         fn=generate_image,
         inputs=[prompt, seed, width, height],
         outputs=[result, seed, latency],
         show_progress="hidden",
-        show_api=False,
+        api_name=False,
         queue=False
     )
 
-    gr.on(
-        triggers=[prompt.submit, prompt.input, width.input, height.input, num_inference_steps.input],
+    generateBtn.click(
         fn=generate_image,
         inputs=[prompt, seed, width, height, randomize_seed, num_inference_steps],
         outputs=[result, seed, latency],
+        show_progress="full",
+        api_name=False,
+        queue=False
+    )
+
+    def update_ui(realtime_enabled):
+        return {
+            prompt: gr.update(interactive=True),
+            generateBtn: gr.update(visible=not realtime_enabled)
+        }
+
+    realtime.change(
+        fn=update_ui,
+        inputs=[realtime],
+        outputs=[prompt, generateBtn],
+        queue=False
+    )
+
+    def realtime_generation(*args):
+        if args[0]:  # If realtime is enabled
+            return next(generate_image(*args[1:]))
+
+    prompt.submit(
+        fn=realtime_generation,
+        inputs=[realtime, prompt, seed, width, height, randomize_seed, num_inference_steps],
+        outputs=[result, seed, latency],
         show_progress="hidden",
-        show_api=False,
-        trigger_mode="always_last",
+        api_name=False,
         queue=False
     )
 
+    for component in [prompt, width, height, num_inference_steps]:
+        component.input(
+            fn=realtime_generation,
+            inputs=[realtime, prompt, seed, width, height, randomize_seed, num_inference_steps],
+            outputs=[result, seed, latency],
+            show_progress="hidden",
+            api_name=False,
+            trigger_mode="always_last",
+            queue=False
+        )
+
 # Launch the app
-demo.launch() 
+demo.launch()