Update app.py

app.py

@@ -1,17 +1,14 @@
 import gradio as gr
-import torch
 from diffusers import AutoPipelineForText2Image
 from io import BytesIO
 from generate_propmts import generate_prompt
 from concurrent.futures import ThreadPoolExecutor, as_completed
-import json
 
 # Load the model once outside of the function
 model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
 
 def generate_image(prompt):
     try:
-        # Truncate prompt if necessary
         output = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
         print(f"Model output: {output}")
 
@@ -21,7 +18,6 @@ def generate_image(prompt):
             buffered = BytesIO()
             image.save(buffered, format="JPEG")
             image_bytes = buffered.getvalue()
-            print(f'image_bytes:{image_bytes}')
             return image_bytes
         else:
             raise Exception("No images returned by the model.")
@@ -31,7 +27,7 @@ def generate_image(prompt):
         return None
 
 def inference(sentence_mapping, character_dict, selected_style):
-    images = {}
+    images = []
     print(f'sentence_mapping: {sentence_mapping}, character_dict: {character_dict}, selected_style: {selected_style}')
     prompts = []
 
@@ -39,23 +35,21 @@ 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((paragraph_number, prompt))
+        prompts.append(prompt)
         print(f"Generated prompt for paragraph {paragraph_number}: {prompt}")
 
     with ThreadPoolExecutor() as executor:
-        future_to_paragraph = {executor.submit(generate_image, prompt): paragraph_number for paragraph_number, prompt in prompts}
+        futures = [executor.submit(generate_image, prompt) for prompt in prompts]
 
-        for future in as_completed(future_to_paragraph):
-            paragraph_number = future_to_paragraph[future]
+        for future in as_completed(futures):
             try:
                 image = future.result()
                 if image:
-                    images[paragraph_number] = image
+                    images.append(image)
             except Exception as e:
-                print(f"Error processing paragraph {paragraph_number}: {e}")
+                print(f"Error processing prompt: {e}")
 
-    # Return the images sorted by paragraph number
-    return [images[paragraph_number] for paragraph_number in sorted(images.keys())]
+    return images
 
 gradio_interface = gr.Interface(
     fn=inference,
@@ -64,7 +58,7 @@ gradio_interface = gr.Interface(
         gr.JSON(label="Character Dict"),
         gr.Dropdown(["oil painting", "sketch", "watercolor"], label="Selected Style")
     ],
-    outputs="json"
+    outputs=gr.Gallery(label="Generated Images")
 )
 
 if __name__ == "__main__":