Update app.py

app.py

@@ -8,7 +8,7 @@ from PIL import Image, ImageDraw, ImageFont
 device = "cuda" if torch.cuda.is_available() else "cpu"
 st.write(f"Using device: {device}")  # Debug message
 
-# Load text model (TinyLlama) with caching
+# Load text model (TinyLlama) with error handling
 @st.cache_resource
 def load_text_model():
     try:
@@ -24,7 +24,7 @@ def load_text_model():
 
 story_generator = load_text_model()
 
-# Load image model (Stable Diffusion) with caching
+# Load image model (Stable Diffusion) with error handling
 @st.cache_resource
 def load_image_model():
     try:
@@ -50,11 +50,12 @@ def generate_story(prompt):
         st.write("⏳ Generating story...")
         story_output = story_generator(
             formatted_prompt,
-            max_length=120,  # Reduced length for efficiency
+            max_length=150,  # Shorter length for efficiency
             do_sample=True,
             temperature=0.7,
             top_k=30,
-            num_return_sequences=1
+            num_return_sequences=1,
+            truncation = True
         )[0]['generated_text']
         st.write("✅ Story generated successfully!")
         return story_output.replace(formatted_prompt, "").strip()
@@ -102,11 +103,10 @@ if user_prompt:
     else:
         with st.spinner("⏳ Generating image..."):
             try:
-                image = image_generator(user_prompt, num_inference_steps=10).images[0]  # Reduced steps from 12 → 10
-                image = image.resize((512, 512))  # Resize to smaller 512x512
+                image = image_generator(user_prompt, num_inference_steps=12).images[0]
                 st.write("✅ Image generated successfully!")
 
-                # Extract first sentence (max 50 characters) for speech bubble
+                # Extract first sentence (50 characters max) for speech bubble
                 speech_text = generated_story.split(".")[0][:50]
                 image_with_bubble = add_speech_bubble(image, speech_text, position=(50, 50))