Update app.py

app.py

@@ -4,36 +4,48 @@ from diffusers import AutoPipelineForText2Image
 import base64
 from io import BytesIO
 from generate_propmts import generate_prompt
+from concurrent.futures import ThreadPoolExecutor
+
 # Load the model once outside of the function
 model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
 
 def generate_image(text, sentence_mapping, character_dict, selected_style):
     try:
-        prompt,_ = generate_prompt(text, sentence_mapping, character_dict, selected_style)
+        prompt, _ = generate_prompt(text, sentence_mapping, character_dict, selected_style)
         image = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
         buffered = BytesIO()
+        image.save(buffered, format="JPEG")
         img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-        if isinstance(result, img_str):
-            image_bytes = base64.b64decode(result)
-            return image_bytes
+        return img_str
     except Exception as e:
+        print(f"Error generating image: {e}")
         return None
 
-def inference(prompt):
-    # Dictionary to store images results
+def inference(text, sentence_mapping, character_dict, selected_style):
     images = {}
-    print(f"Received grouped_sentences: {grouped_sentences}")
-    # Debugging statement
-    with concurrent.images.ThreadPoolExecutor() as executor:
+    # Here we assume `sentence_mapping` is a dictionary where keys are paragraph numbers and values are lists of sentences
+    grouped_sentences = sentence_mapping
+
+    with ThreadPoolExecutor() as executor:
+        futures = {}
         for paragraph_number, sentences in grouped_sentences.items():
             combined_sentence = " ".join(sentences)
-            images[paragraph_number] = executor.submit(generate_image, combined_sentence, sentence_mapping, general_descriptions, selected_style)
+            futures[paragraph_number] = executor.submit(generate_image, combined_sentence, sentence_mapping, character_dict, selected_style)
+
+        for paragraph_number, future in futures.items():
+            images[paragraph_number] = future.result()
+
     return images
 
 gradio_interface = gr.Interface(
     fn=inference,
-    inputs="text",
-    outputs="text"  # Change output to text to return base64 string
+    inputs=[
+        gr.inputs.Textbox(label="Text"),
+        gr.inputs.Textbox(label="Sentence Mapping"),
+        gr.inputs.Textbox(label="Character Dict"),
+        gr.inputs.Dropdown(["Style 1", "Style 2", "Style 3"], label="Selected Style")
+    ],
+    outputs="json"  # Return the dictionary of images
 )
 
 if __name__ == "__main__":