add example selection and processing pipeline functions

app.py

@@ -22,7 +22,7 @@ _image_gen_pipeline = None
 @spaces.GPU()
 def get_image_gen_pipeline():
     global _image_gen_pipeline
-    if _image_gen_pipeline is None:
+    if (_image_gen_pipeline is None):
         try:
             device = "cuda" if torch.cuda.is_available() else "cpu"
             dtype = torch.bfloat16
@@ -42,7 +42,7 @@ def get_image_gen_pipeline():
 @spaces.GPU()
 def get_text_gen_pipeline():
     global _text_gen_pipeline
-    if _text_gen_pipeline is None:
+    if (_text_gen_pipeline is None):
         try:
             device = "cuda" if torch.cuda.is_available() else "cpu"
             tokenizer = AutoTokenizer.from_pretrained(
@@ -179,6 +179,36 @@ def preload_models():
     print(status)
     return success
 
+# Add a new function to handle example selection
+def handle_example_click(example_prompt):
+    # Immediately return the example prompt to update the UI
+    return example_prompt, "Example selected - click 'Refine prompt with Mistral' to process"
+
+# Create a combined function that handles the whole pipeline from example to image
+@spaces.GPU()
+def process_example_pipeline(example_prompt, seed, randomize_seed, width, height, num_inference_steps, progress=gr.Progress()):
+    # Step 1: Update status
+    progress(0, desc="Starting example processing")
+    progress_status = "Selected example: " + example_prompt
+    
+    # Step 2: Refine the prompt
+    progress(0.1, desc="Refining prompt with Mistral")
+    refined, status = refine_prompt(example_prompt, progress)
+    
+    if not refined:
+        return example_prompt, "", None, "Failed to refine prompt: " + status
+    
+    progress(0.5, desc="Prompt refined, generating image")
+    progress_status = status
+    
+    # Step 3: Generate the image
+    image, image_status = infer(refined, seed, randomize_seed, width, height, num_inference_steps, progress)
+    
+    progress(1.0, desc="Process complete")
+    final_status = f"{progress_status} → {image_status}"
+    
+    return example_prompt, refined, image, final_status
+
 def create_interface():
     # Preload models if needed
     if PRELOAD_MODELS:
@@ -246,16 +276,28 @@ def create_interface():
                         minimum=1,
                         maximum=50,
                         step=1,
-                        value=10,
+                        value=6,
                     )
 
-            # Examples section
+            # Examples section - use the pipeline function for examples
             gr.Examples(
                 examples=examples,
-                fn=refine_prompt,
-                inputs=[prompt],
-                outputs=[refined_prompt],
-                cache_examples=True,
+                fn=process_example_pipeline,
+                inputs=[
+                    prompt,
+                    seed,
+                    randomize_seed, 
+                    width, 
+                    height,
+                    num_inference_steps
+                ],
+                outputs=[
+                    prompt,
+                    refined_prompt,
+                    generated_image, 
+                    error_box
+                ],
+                cache_examples=True,  # Can be cached now
             )
 
         # Event handlers