Update app.py

app.py

@@ -1,71 +1,36 @@
 import gradio as gr
+import torch
 from diffusers import AutoPipelineForText2Image
-from diffusers.schedulers import DPMSolverMultistepScheduler
-from generate_propmts import generate_prompt  # Assuming you have this module
-from PIL import Image
+from io import BytesIO
+from generate_propmts import generate_prompt
+from concurrent.futures import ThreadPoolExecutor
 import asyncio
-import threading
-import traceback
 
-# Define the SchedulerWrapper class
-class SchedulerWrapper:
-    def __init__(self, scheduler):
-        self.scheduler = scheduler
-        self._step = threading.local()
-        self._step.step = 0
-
-    def __getattr__(self, name):
-        return getattr(self.scheduler, name)
-
-    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)
-
-    @property
-    def timesteps(self):
-        return self.scheduler.timesteps
-
-    def set_timesteps(self, *args, **kwargs):
-        return self.scheduler.set_timesteps(*args, **kwargs)
-
-# Load the model and wrap the scheduler
+# Load the model once outside of the function
 model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
 
-scheduler = DPMSolverMultistepScheduler.from_config(model.scheduler.config)
-wrapped_scheduler = SchedulerWrapper(scheduler)
-model.scheduler = wrapped_scheduler
-
-# Define the image generation function
 async def generate_image(prompt):
     try:
-        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,  # 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]
+        # Truncate prompt if necessary
+        output = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
+        print(f"Model output: {output}")
+
+        # Check if the model returned images
+        if isinstance(output.images, list) and len(output.images) > 0:
+            image = output.images[0]
+            buffered = BytesIO()
+            image.save(buffered, format="JPEG")
+            image_bytes = buffered.getvalue()
+            return image_bytes
         else:
             raise Exception("No images returned by the model.")
+            
     except Exception as e:
         print(f"Error generating image: {e}")
-        traceback.print_exc()
-        return None  # Return None on error to handle it gracefully in the UI
+        return None
 
-# Define the inference function
-async def inference(sentence_mapping, character_dict, selected_style):
-    images = []
+async def process_prompt(sentence_mapping, character_dict, selected_style):
+    images = {}
     print(f'sentence_mapping: {sentence_mapping}, character_dict: {character_dict}, selected_style: {selected_style}')
     prompts = []
 
@@ -73,29 +38,36 @@ async def inference(sentence_mapping, character_dict, selected_style):
     for paragraph_number, sentences in sentence_mapping.items():
         combined_sentence = " ".join(sentences)
         prompt = generate_prompt(combined_sentence, sentence_mapping, character_dict, selected_style)
-        prompts.append(prompt)
+        prompts.append((paragraph_number, prompt))
         print(f"Generated prompt for paragraph {paragraph_number}: {prompt}")
 
-    # Use asyncio.gather to run generate_image in parallel
-    tasks = [generate_image(prompt) for prompt in prompts]
-    images = await asyncio.gather(*tasks)
-
-    # Filter out None values
-    images = [image for image in images if image is not None]
+    loop = asyncio.get_event_loop()
+    tasks = []
+    
+    with ThreadPoolExecutor() as executor:
+        for paragraph_number, prompt in prompts:
+            tasks.append(loop.run_in_executor(executor, generate_image, prompt))
+
+        for paragraph_number, task in zip(sentence_mapping.keys(), await asyncio.gather(*tasks)):
+            try:
+                image = task
+                if image:
+                    images[paragraph_number] = image
+            except Exception as e:
+                print(f"Error processing paragraph {paragraph_number}: {e}")
 
     return images
 
-# Define the Gradio interface
 gradio_interface = gr.Interface(
-    fn=inference,
+    fn=process_prompt,
     inputs=[
         gr.JSON(label="Sentence Mapping"),
         gr.JSON(label="Character Dict"),
         gr.Dropdown(["oil painting", "sketch", "watercolor"], label="Selected Style")
     ],
-    outputs=gr.Gallery(label="Generated Images")
+    outputs="json",
+    concurrency_limit=10  # Allow up to 10 concurrent executions
 )
 
-# Run the Gradio app
 if __name__ == "__main__":
     gradio_interface.launch()