Update app.py

app.py

@@ -11,19 +11,19 @@ model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
 
 def generate_image(prompt):
     try:
-        # Generate the image
+        # Truncate prompt if necessary
         output = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
         print(f"Model output: {output}")
 
-        # The model returns a PIL image
-        if output:
-            image = output[0]  # Extract the PIL image from the output
+        # Check if the model returned images
+        if isinstance(output.images, list) and len(output.images) > 0:
+            image = output.images[0]
             buffered = BytesIO()
             image.save(buffered, format="JPEG")
             image_bytes = buffered.getvalue()
             return image_bytes
         else:
-            raise Exception("No image returned by the model.")
+            raise Exception("No images returned by the model.")
             
     except Exception as e:
         print(f"Error generating image: {e}")
@@ -37,7 +37,7 @@ def inference(sentence_mapping, character_dict, selected_style):
     # 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)
+        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}")