Update app.py

app.py

@@ -1,66 +1,28 @@
-import threading
+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 traceback
 
-class Scheduler:
-    def __init__(self):
-        self._step = threading.local()
-        self._step.step = None
-
-    def _init_step_index(self):
-        self._step.step = 0
-
-    @property
-    def step_index(self):
-        if self._step.step is None:
-            self._init_step_index()
-        return self._step.step
-
-    @step_index.setter
-    def step_index(self, value):
-        self._step.step = value
-
-    def step_process(self, noise_pred, t, latents, **extra_step_kwargs):
-        try:
-            sigma_to = self.sigmas[self.step_index + 1]
-            self.step_index += 1
-            # Process the step (pseudocode)
-            # latents = process_latents(noise_pred, t, latents, sigma_to, **extra_step_kwargs)
-            return latents
-        except IndexError as e:
-            print(f"Index error during step processing: {e}")
-            traceback.print_exc()
-            return latents
-
-# Mocking a model class for demonstration purposes
-class MockModel:
-    def __init__(self):
-        self.scheduler = Scheduler()
-
-    def __call__(self, prompt, num_inference_steps, guidance_scale):
-        # Simulate the inference steps
-        latents = None
-        for t in range(num_inference_steps):
-            noise_pred = None  # Replace with actual noise prediction
-            latents = self.scheduler.step_process(noise_pred, t, latents)
-        return {"images": [Image.new("RGB", (512, 512))]}  # Return a dummy image for now
-
-# Load the actual model
-model = MockModel()
+# Load the model once outside of the function
+model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
 
 def generate_image(prompt):
     try:
-        output = model(prompt=prompt, num_inference_steps=3, guidance_scale=0.0)
+        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:
-            return output['images'][0]
+        if isinstance(output.images, list) and len(output.images) > 0:
+            return output.images[0]
         else:
             raise Exception("No images returned by the model.")
             
+    except IndexError as e:
+        print(f"Index error during image generation: {e}")
+        traceback.print_exc()
+        return None
     except Exception as e:
         print(f"Error generating image: {e}")
         traceback.print_exc()
@@ -71,6 +33,7 @@ def inference(sentence_mapping, character_dict, selected_style):
     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)
@@ -99,7 +62,7 @@ gradio_interface = gr.Interface(
         gr.Dropdown(["oil painting", "sketch", "watercolor"], label="Selected Style")
     ],
     outputs=gr.Gallery(label="Generated Images")
-)
+.queue(default_concurrency_limit=5)
 
 if __name__ == "__main__":
     gradio_interface.launch()