Update app.py

app.py

@@ -3,27 +3,19 @@ import torch
 from diffusers import AutoPipelineForText2Image
 from io import BytesIO
 from generate_propmts import generate_prompt
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, as_completed
 import json
 
 # Load the model once outside of the function
 model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
 
-# Helper function to truncate prompt to fit the model's maximum sequence length
-# def truncate_prompt(prompt, max_length=77):
-#     return prompt[:max_length]
 
-def generate_image(text, sentence_mapping, character_dict, selected_style):
+def generate_image(prompt):
     try:
-        prompt = generate_prompt(text, sentence_mapping, character_dict, selected_style)
-        print(f"Generated prompt: {prompt}")
         # Truncate prompt if necessary
-        # prompt = truncate_prompt(prompt)
-        # print(f"truncate_prompt: {prompt}")
-        output = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
+        output = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
         print(f"Model output: {output}")
-        print("len of output:", len(output))
-        print("output.images:", output.images)
+
         # Check if the model returned images
         if output.images:
             image = output.images[0]
@@ -41,18 +33,27 @@ def generate_image(text, sentence_mapping, character_dict, selected_style):
 def inference(sentence_mapping, character_dict, selected_style):
     images = {}
     print(f'sentence_mapping: {sentence_mapping}, character_dict: {character_dict}, selected_style: {selected_style}')
-    # Here we assume `sentence_mapping` is a dictionary where keys are paragraph numbers and values are lists of sentences
-    grouped_sentences = sentence_mapping
+    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((paragraph_number, prompt))
+        print(f"Generated prompt for paragraph {paragraph_number}: {prompt}")
 
     with ThreadPoolExecutor() as executor:
-        futures = {}
-        for paragraph_number, sentences in grouped_sentences.items():
-            combined_sentence = " ".join(sentences)
-            futures[paragraph_number] = executor.submit(generate_image, combined_sentence, sentence_mapping, character_dict, selected_style)
+        future_to_paragraph = {executor.submit(generate_image, prompt): paragraph_number for paragraph_number, prompt in prompts}
+
+        for future in as_completed(future_to_paragraph):
+            paragraph_number = future_to_paragraph[future]
+            try:
+                image = future.result()
+                if image:
+                    images[paragraph_number] = image
+            except Exception as e:
+                print(f"Error processing paragraph {paragraph_number}: {e}")
 
-    for paragraph_number, future in futures.items():
-        images[paragraph_number] = future.result()
-    print(f'images:{images}')
     return images
 
 gradio_interface = gr.Interface(
@@ -60,7 +61,7 @@ gradio_interface = gr.Interface(
     inputs=[
         gr.JSON(label="Sentence Mapping"),
         gr.JSON(label="Character Dict"),
-        gr.Dropdown(["Style 1", "Style 2", "Style 3"], label="Selected Style")
+        gr.Dropdown(["oil painting", "sketch", "watercolor"], label="Selected Style")
     ],
     outputs="json"
 )