Update app.py

app.py

@@ -9,22 +9,21 @@ import json
 # Load the model once outside of the function
 model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
 
-
 def generate_image(prompt):
     try:
-        # Truncate prompt if necessary
-        output = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0).images[0]
+        # Generate the image
+        output = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
         print(f"Model output: {output}")
 
-        # Check if the model returned images
-        if output.images:
-            image = output.images[0]
+        # The model returns a PIL image
+        if output:
+            image = output[0]  # Extract the PIL image from the output
             buffered = BytesIO()
             image.save(buffered, format="JPEG")
             image_bytes = buffered.getvalue()
             return image_bytes
         else:
-            raise Exception("No images returned by the model.")
+            raise Exception("No image returned by the model.")
             
     except Exception as e:
         print(f"Error generating image: {e}")
@@ -38,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}")