Update app.py

app.py

@@ -9,16 +9,33 @@ 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):
+    tokens = prompt.split()
+    if len(tokens) > max_length:
+        return ' '.join(tokens[:max_length])
+    return prompt
+
 def generate_image(text, sentence_mapping, character_dict, selected_style):
     try:
         prompt, _ = generate_prompt(text, sentence_mapping, character_dict, selected_style)
         print(f"Generated prompt: {prompt}")
+        
+        # Truncate prompt if necessary
+        prompt = truncate_prompt(prompt)
         output = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
         print(f"Model output: {output}")
-        image = output.images[0]
-        buffered = BytesIO()
-        image_bytes = buffered.getvalue()
-        return image_bytes
+
+        # Check if the model returned images
+        if output.images:
+            image = output.images[0]
+            buffered = BytesIO()
+            image.save(buffered, format="JPEG")
+            image_bytes = buffered.getvalue()
+            return image_bytes
+        else:
+            raise Exception("No images returned by the model.")
+            
     except Exception as e:
         print(f"Error generating image: {e}")
         return None
@@ -47,7 +64,7 @@ gradio_interface = gr.Interface(
         gr.JSON(label="Character Dict"),
         gr.Dropdown(["Style 1", "Style 2", "Style 3"], label="Selected Style")
     ],
-    outputs="json"  
+    outputs="json"
 )
 
 if __name__ == "__main__":