Update app.py

app.py

@@ -2,42 +2,51 @@ import gradio as gr
 import threading
 import os
 import torch
+from diffusers import StableDiffusionPipeline
 
+# Set environment variables for performance
 os.environ["OMP_NUM_THREADS"] = str(os.cpu_count())
 torch.set_num_threads(os.cpu_count())
 
-model1 = gr.load("models/prithivMLmods/SD3.5-Turbo-Realism-2.0-LoRA")
-model2 = gr.load("models/Purz/face-projection")
+# Load models
+model1 = StableDiffusionPipeline.from_pretrained("prithivMLmods/SD3.5-Turbo-Realism-2.0-LoRA")
+model2 = StableDiffusionPipeline.from_pretrained("Purz/face-projection")
 
+# Event to stop image generation
 stop_event = threading.Event()
 
 def generate_images(text, selected_model):
+    """Generate images based on the input text and selected model."""
     stop_event.clear()
 
-    if selected_model == "Model 1 (Turbo Realism)":
-        model = model1
-    elif selected_model == "Model 2 (Face Projection)":
-        model = model2
-    else:
-        return ["Invalid model selection."] * 3
+    try:
+        if selected_model == "Model 1 (Turbo Realism)":
+            model = model1
+        elif selected_model == "Model 2 (Face Projection)":
+            model = model2
+        else:
+            return ["Invalid model selection."] * 3
 
-    results = []
-    for i in range(3):
-        if stop_event.is_set():
-            return ["Image generation stopped by user."] * 3
+        results = []
+        for i in range(3):
+            if stop_event.is_set():
+                return ["Image generation stopped by user."] * 3
 
-        modified_text = f"{text} variation {i+1}"
-        result = model(modified_text)
-        results.append(result)
+            modified_text = f"{text} variation {i+1}"
+            result = model(modified_text).images[0]  # Generate image
+            results.append(result)
 
-    return results
+        return results
+    except Exception as e:
+        return [f"Error generating images: {str(e)}"] * 3
 
 def stop_generation():
-    """Stops the ongoing image generation by setting the stop_event flag."""
+    """Stop the ongoing image generation."""
     stop_event.set()
     return ["Generation stopped."] * 3
 
-with gr.Blocks() as interface:#...
+# Gradio interface
+with gr.Blocks() as interface:
     gr.Markdown(
         "### ⚠ Sorry for the inconvenience. The Space is currently running on the CPU, which might affect performance. We appreciate your understanding."
     )
@@ -61,4 +70,5 @@ with gr.Blocks() as interface:#...
     generate_button.click(generate_images, inputs=[text_input, model_selector], outputs=[output1, output2, output3])
     stop_button.click(stop_generation, inputs=[], outputs=[output1, output2, output3])
 
-interface.launch()
+# Launch the interface with an API endpoint
+interface.launch(api_name="generate_images")