Update app.py

app.py

@@ -1,62 +1,88 @@
 import gradio as gr
 from diffusers import AutoPipelineForText2Image
-from transformers import AutoTokenizer
+from diffusers.schedulers import DPMSolverMultistepScheduler
+from generate_propmts import generate_prompt  # Assuming you have this module
 from PIL import Image
 import asyncio
+import threading
+import traceback
 
+# Define the SchedulerWrapper class
 class SchedulerWrapper:
     def __init__(self, scheduler):
         self.scheduler = scheduler
-        
-    def __getattr__(self, name):
-        return getattr(self.scheduler, name)
-        
+        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)
+
     @property
     def timesteps(self):
         return self.scheduler.timesteps
 
-    def set_timesteps(self, timesteps):
-        self.scheduler.set_timesteps(timesteps)
+    def set_timesteps(self, *args, **kwargs):
+        return self.scheduler.set_timesteps(*args, **kwargs)
 
-# Load the model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("stabilityai/sdxl-turbo")
+# Load the model and wrap the scheduler
 model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
 
-# Wrap the scheduler
-scheduler = model.scheduler
+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
+        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"
+            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:
             raise Exception("No images returned by the model.")
     except Exception as e:
         print(f"Error generating image: {e}")
-        return None
+        traceback.print_exc()
+        return None  # Return None on error to handle it gracefully in the UI
 
+# Define the inference function
 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 = []
+
+    # Generate prompts for each paragraph
     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)
+        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]
+
     return images
 
+# Define the Gradio interface
 gradio_interface = gr.Interface(
     fn=inference,
     inputs=[
@@ -67,5 +93,6 @@ gradio_interface = gr.Interface(
     outputs=gr.Gallery(label="Generated Images")
 )
 
+# Run the Gradio app
 if __name__ == "__main__":
     gradio_interface.launch()