Update app.py

app.py

@@ -1,3 +1,4 @@
+import json
 import gradio as gr
 import torch
 from diffusers import AutoPipelineForText2Image
@@ -14,10 +15,9 @@ def generate_image(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()
+        buffered.write(image.tobytes())
         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
@@ -25,14 +25,19 @@ 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
+    
+    # Parse sentence_mapping JSON string into a dictionary
+    try:
+        grouped_sentences = json.loads(sentence_mapping)
+    except json.JSONDecodeError as e:
+        print(f"Error parsing JSON: {e}")
+        return {"error": "Invalid JSON input for sentence_mapping"}
 
     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)
+            futures[paragraph_number] = executor.submit(generate_image, combined_sentence, grouped_sentences, character_dict, selected_style)
 
         for paragraph_number, future in futures.items():
             images[paragraph_number] = future.result()
@@ -42,11 +47,11 @@ def inference(sentence_mapping, character_dict, selected_style):
 gradio_interface = gr.Interface(
     fn=inference,
     inputs=[
-        gr.Textbox(label="Sentence Mapping"),
-        gr.Textbox(label="Character Dict"),
+        gr.Textbox(label="Sentence Mapping (JSON)"),
+        gr.Textbox(label="Character Dict (JSON)"),
         gr.Dropdown(["Style 1", "Style 2", "Style 3"], label="Selected Style")
     ],
-    outputs="text"  
+    outputs="text"
 )
 
 if __name__ == "__main__":