Update app.py

app.py

@@ -1,4 +1,3 @@
-import json
 import gradio as gr
 import torch
 from diffusers import AutoPipelineForText2Image
@@ -6,6 +5,7 @@ import base64
 from io import BytesIO
 from generate_propmts import generate_prompt
 from concurrent.futures import ThreadPoolExecutor
+import json
 
 # Load the model once outside of the function
 model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
@@ -15,9 +15,10 @@ 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")
-        return img_str
+        if isinstance(result, img_str):
+            image_bytes = base64.b64decode(result)
+        return image_bytes
     except Exception as e:
         print(f"Error generating image: {e}")
         return None
@@ -25,19 +26,14 @@ 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}')
-    
-    # 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"}
+    # 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)
-            futures[paragraph_number] = executor.submit(generate_image, combined_sentence, grouped_sentences, character_dict, 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()
@@ -47,11 +43,11 @@ def inference(sentence_mapping, character_dict, selected_style):
 gradio_interface = gr.Interface(
     fn=inference,
     inputs=[
-        gr.Textbox(label="Sentence Mapping (JSON)"),
-        gr.Textbox(label="Character Dict (JSON)"),
+        gr.JSON(label="Sentence Mapping"),
+        gr.JSON(label="Character Dict"),
         gr.Dropdown(["Style 1", "Style 2", "Style 3"], label="Selected Style")
     ],
-    outputs="text"
+    outputs="json"  
 )
 
 if __name__ == "__main__":