Update app.py

app.py

@@ -1,79 +1,82 @@
 import os
 from io import BytesIO
+from PIL import Image
 from diffusers import AutoPipelineForText2Image
 import gradio as gr
 import base64
-from generate_prompts import generate_prompt
 
 # Load the model once at the start
 print("Loading the Stable Diffusion model...")
-model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
-print("Model loaded successfully.")
-
-def truncate_prompt(prompt, max_length=77):
-    tokens = prompt.split()
-    if len(tokens) > max_length:
-        prompt = " ".join(tokens[:max_length])
-    return prompt
+try:
+    model = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo")
+    print("Model loaded successfully.")
+except Exception as e:
+    print(f"Error loading model: {e}")
+    model = None
 
 def generate_image(prompt):
     try:
-        truncated_prompt = truncate_prompt(prompt)
-        print(f"Generating image with truncated prompt: {truncated_prompt}")
-        
-        # Call the model
-        output = model(prompt=truncated_prompt, num_inference_steps=1, guidance_scale=0.0)
+        if model is None:
+            raise ValueError("Model not loaded properly.")
         
-        # Check if output is valid
-        if output is not None and hasattr(output, 'images') and output.images:
-            print(f"Image generated")
+        print(f"Generating image with prompt: {prompt}")
+        output = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
+        print(f"Model output: {output}")
+
+        if output is None:
+            raise ValueError("Model returned None")
+
+        if hasattr(output, 'images') and output.images:
+            print(f"Image generated successfully")
             image = output.images[0]
-            return image, None
+            buffered = BytesIO()
+            image.save(buffered, format="JPEG")
+            image_bytes = buffered.getvalue()
+            img_str = base64.b64encode(image_bytes).decode("utf-8")
+            print("Image encoded to base64")
+            print(f'img_str: {img_str[:100]}...')  # Print a snippet of the base64 string
+            return img_str, None
         else:
-            print(f"No images found or generated output is None")
-            return None, "No images found or generated output is None"
+            print(f"No images found in model output")
+            raise ValueError("No images found in model output")
     except Exception as e:
         print(f"An error occurred while generating image: {e}")
         return None, str(e)
 
-def inference(prompt):
-    print(f"Received prompt: {prompt}")  # Debugging statement
-    image, error = generate_image(prompt)
-    if error:
-        print(f"Error generating image: {error}")  # Debugging statement
-        return "Error: " + error
-
-    buffered = BytesIO()
-    image.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-    return img_str
-
-def process_prompt(sentence_mapping, character_dict, selected_style):
-    print("Processing prompt...")
-    print(f"Sentence Mapping: {sentence_mapping}")
-    print(f"Character Dict: {character_dict}")
-    print(f"Selected Style: {selected_style}")
+def inference(sentence_mapping, character_dict, selected_style):
+    try:
+        print(f"Received sentence_mapping: {sentence_mapping}, type: {type(sentence_mapping)}")
+        print(f"Received character_dict: {character_dict}, type: {type(character_dict)}")
+        print(f"Received selected_style: {selected_style}, type: {type(selected_style)}")
 
-    images = {}
-    for paragraph_number, sentences in sentence_mapping.items():
-        combined_sentence = " ".join(sentences)
-        prompt, _ = generate_prompt(combined_sentence, sentence_mapping, character_dict, selected_style)
-        img_str = inference(prompt)
-        images[paragraph_number] = img_str
+        if sentence_mapping is None or character_dict is None or selected_style is None:
+            return {"error": "One or more inputs are None"}
 
-    print("Prompt processing complete. Generated images: ", images)
-    return images
+        images = {}
+        for paragraph_number, sentences in sentence_mapping.items():
+            combined_sentence = " ".join(sentences)
+            prompt = f"Make an illustration in {selected_style} style from: {combined_sentence}"
+            print(f"Generated prompt for paragraph {paragraph_number}: {prompt}")
+            img_str, error = generate_image(prompt)
+            if error:
+                images[paragraph_number] = f"Error: {error}"
+            else:
+                images[paragraph_number] = img_str
+        return images
+    except Exception as e:
+        print(f"An error occurred during inference: {e}")
+        return {"error": str(e)}
 
 gradio_interface = gr.Interface(
-    fn=process_prompt,
+    fn=inference,
     inputs=[
         gr.JSON(label="Sentence Mapping"),
         gr.JSON(label="Character Dict"),
         gr.Dropdown(["oil painting", "sketch", "watercolor"], label="Selected Style")
     ],
     outputs="json"
-).queue(default_concurrency_limit=20)  # Set concurrency limit if needed
+)
 
 if __name__ == "__main__":
     print("Launching Gradio interface...")
-    gradio_interface.launch()
+    gradio_interface.launch()