Update app.py

app.py

@@ -1,37 +1,27 @@
 import gradio as gr
 from diffusers import AutoPipelineForText2Image
 from generate_propmts import generate_prompt
-from concurrent.futures import ThreadPoolExecutor, as_completed
 from PIL import Image
-import threading
+import asyncio
 import traceback
 
 # Load the model once outside of the function
 model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
 
-# Create a thread-local storage for step indices
-scheduler_step_storage = threading.local()
-
-def generate_image(prompt):
+async def generate_image(prompt):
     try:
-        # Initialize step index per thread if not already set
-        if not hasattr(scheduler_step_storage, 'step'):
-            scheduler_step_storage.step = 0
-        
         # Use a sensible default for num_inference_steps
         num_inference_steps = 1  # Adjust this value as needed
-        
-        # Use the thread-local step index
-        output = model(
-            prompt=prompt, 
+
+        # Use the model to generate an image
+        output = await asyncio.to_thread(
+            model,
+            prompt=prompt,
             num_inference_steps=num_inference_steps,  # Use a higher value for inference steps
             guidance_scale=0.0,  # Typical value for guidance scale in image generation
             output_type="pil"  # Directly get PIL Image objects
         )
 
-        # Increment the step index after generating the image
-        scheduler_step_storage.step += 1  
-        
         # Check for output validity and return
         if output.images:
             return output.images[0]
@@ -43,7 +33,7 @@ def generate_image(prompt):
         traceback.print_exc()
         return None  # Return None on error to handle it gracefully in the UI
 
-def inference(sentence_mapping, character_dict, selected_style):
+async def inference(sentence_mapping, character_dict, selected_style):
     images = []
     print(f'sentence_mapping: {sentence_mapping}, character_dict: {character_dict}, selected_style: {selected_style}')
     prompts = []
@@ -55,17 +45,12 @@ def inference(sentence_mapping, character_dict, selected_style):
         prompts.append(prompt)
         print(f"Generated prompt for paragraph {paragraph_number}: {prompt}")
 
-    with ThreadPoolExecutor() as executor:
-        futures = [executor.submit(generate_image, prompt) for prompt in prompts]
+    # Use asyncio.gather to run generate_image in parallel
+    tasks = [generate_image(prompt) for prompt in prompts]
+    images = await asyncio.gather(*tasks)
 
-        for future in as_completed(futures):
-            try:
-                image = future.result()
-                if image:
-                    images.append(image)
-            except Exception as e:
-                print(f"Error processing prompt: {e}")
-                traceback.print_exc()
+    # Filter out None values
+    images = [image for image in images if image is not None]
 
     return images