Update app.py

app.py

@@ -1,25 +1,46 @@
 import gradio as gr
-from diffusers import AutoPipelineForText2Image, EulerAncestralDiscreteScheduler
+from diffusers import AutoPipelineForText2Image
 from generate_propmts import generate_prompt
 from PIL import Image
 import asyncio
+import threading
 import traceback
 
-# Load the model with a different scheduler
+# Load the model once outside of the function
 model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
-model.scheduler = DPMSolverMultistepScheduler.from_config(model.scheduler.config)  # Changed scheduler
 
+class SchedulerWrapper:
+    def __init__(self, scheduler):
+        self.scheduler = scheduler
+        self._step = threading.local()
+        self._step.step = 0
+
+    def step(self, *args, **kwargs):
+        try:
+            self._step.step += 1
+            return self.scheduler.step(*args, **kwargs)
+        except IndexError:
+            self._step.step = 0
+            return self.scheduler.step(*args, **kwargs)
+
+# Wrap the scheduler
+model.scheduler = SchedulerWrapper(model.scheduler)
 
 async def generate_image(prompt):
     try:
-        num_inference_steps = 5  # You can adjust this
+        # Set a higher value for num_inference_steps
+        num_inference_steps = 5  # Adjust this value as needed
+
+        # Use the model to generate an image
         output = await asyncio.to_thread(
             model,
             prompt=prompt,
-            num_inference_steps=num_inference_steps, 
-            guidance_scale=0.0,
-            output_type="pil"
+            num_inference_steps=num_inference_steps,
+            guidance_scale=0.0,  # Typical value for guidance scale in image generation
+            output_type="pil"  # Directly get PIL Image objects
         )
+
+        # Check for output validity and return
         if output.images:
             return output.images[0]
         else:
@@ -27,8 +48,7 @@ async def generate_image(prompt):
     except Exception as e:
         print(f"Error generating image: {e}")
         traceback.print_exc()
-        return None
- 
+        return None  # Return None on error to handle it gracefully in the UI
 
 async def inference(sentence_mapping, character_dict, selected_style):
     images = []