Update app.py

app.py

@@ -2,29 +2,28 @@ import os
 import asyncio
 from concurrent.futures import ProcessPoolExecutor
 from io import BytesIO
-from PIL import Image
 from diffusers import StableDiffusionPipeline
 import gradio as gr
 from generate_prompts import generate_prompt
 
 # Load the model once at the start
 print("Loading the Stable Diffusion model...")
-try:
-    model = StableDiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo")
-    print("Model loaded successfully.")
-except Exception as e:
-    print(f"Error loading model: {e}")
-    model = None
+model = StableDiffusionPipeline.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
 
 def generate_image(prompt, prompt_name):
     try:
-        if model is None:
-            raise ValueError("Model not loaded properly.")
-        
         print(f"Generating image for {prompt_name} with prompt: {prompt}")
-        output = model(prompt=prompt, num_inference_steps=1, guidance_scale=0.0)
+        truncated_prompt = truncate_prompt(prompt)
+        output = model(prompt=truncated_prompt, num_inference_steps=1, guidance_scale=0.0)
         print(f"Model output for {prompt_name}: {output}")
-
+        
         if output and hasattr(output, 'images') and output.images:
             print(f"Image generated for {prompt_name}")
             image = output.images[0]