Update app.py

app.py

@@ -1,28 +1,66 @@
-import gradio as gr
+import threading
 from diffusers import AutoPipelineForText2Image
-from generate_propmts import generate_prompt
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from PIL import Image
 import traceback
 
-# Load the model once outside of the function
-model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
+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()
 
 def generate_image(prompt):
     try:
-        output = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
+        output = model(prompt=prompt, num_inference_steps=3, 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()
@@ -33,7 +71,6 @@ 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)