Update app.py

app.py

@@ -4,43 +4,65 @@ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 from diffusers import StableDiffusionPipeline
 from PIL import Image, ImageDraw, ImageFont
 
-# Check if CUDA is available for GPU acceleration
+# Check for GPU availability
 device = "cuda" if torch.cuda.is_available() else "cpu"
+st.write(f"Using device: {device}")  # Debug message
 
-# Load the text generation model (TinyLlama)
+# Load text model (TinyLlama) with error handling
 @st.cache_resource
 def load_text_model():
-    st.write("Loading text model...")  # Debug message
-    model_name = "distilgpt2"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
-    st.write("Model loaded!")  # Confirm model is ready
-    return pipeline("text-generation", model=model, tokenizer=tokenizer)
+    try:
+        st.write("⏳ Loading text model...")
+        model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
+        st.write("✅ Text model loaded successfully!")
+        return pipeline("text-generation", model=model, tokenizer=tokenizer)
+    except Exception as e:
+        st.error(f"❌ Error loading text model: {e}")
+        return None
 
 story_generator = load_text_model()
 
-# Load the image generation model (Stable Diffusion Turbo)
+# Load image model (Stable Diffusion) with error handling
 @st.cache_resource
 def load_image_model():
-    model_id = "runwayml/stable-diffusion-v1-5"
-    return StableDiffusionPipeline.from_pretrained(model_id).to(device)
+    try:
+        st.write("⏳ Loading image model...")
+        model_id = "runwayml/stable-diffusion-v1-5"
+        model = StableDiffusionPipeline.from_pretrained(model_id).to(device)
+        st.write("✅ Image model loaded successfully!")
+        return model
+    except Exception as e:
+        st.error(f"❌ Error loading image model: {e}")
+        return None
 
 image_generator = load_image_model()
 
 # Function to generate a short story
 def generate_story(prompt):
+    if not story_generator:
+        return "❌ Error: Story model not loaded."
+
     formatted_prompt = f"Write a short comic-style story about: {prompt}\n\nStory:"
-    story_output = story_generator(
-        formatted_prompt,
-        max_length=150,  # Reduce text generation length
-        do_sample=True,
-        temperature=0.7,
-        top_k=30,  # Lower value for efficiency
-        num_return_sequences=1
-    )[0]['generated_text']
-    return story_output.replace(formatted_prompt, "").strip()
-
-# Function to add a speech bubble to the image
+    
+    try:
+        st.write("⏳ Generating story...")
+        story_output = story_generator(
+            formatted_prompt,
+            max_length=150,  # Shorter length for efficiency
+            do_sample=True,
+            temperature=0.7,
+            top_k=30,
+            num_return_sequences=1
+        )[0]['generated_text']
+        st.write("✅ Story generated successfully!")
+        return story_output.replace(formatted_prompt, "").strip()
+    except Exception as e:
+        st.error(f"❌ Error generating story: {e}")
+        return "Error generating story."
+
+# Function to add a speech bubble to an image
 def add_speech_bubble(image, text, position=(50, 50)):
     draw = ImageDraw.Draw(image)
 
@@ -74,10 +96,19 @@ if user_prompt:
     st.write(generated_story)
 
     st.subheader("🖼️ AI-Generated Image")
-    with st.spinner("Generating image..."):
-        image = image_generator(user_prompt, num_inference_steps=30).images[0]
-
-    speech_text = generated_story.split(".")[0][:50]
-    image_with_bubble = add_speech_bubble(image, speech_text, position=(50, 50))
-
-    st.image(image_with_bubble, caption="Generated Comic Image", use_container_width=True)
+    
+    if not image_generator:
+        st.error("❌ Error: Image model not loaded.")
+    else:
+        with st.spinner("⏳ Generating image..."):
+            try:
+                image = image_generator(user_prompt, num_inference_steps=30).images[0]
+                st.write("✅ Image generated successfully!")
+
+                # 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))
+
+                st.image(image_with_bubble, caption="Generated Comic Image", use_container_width=True)
+            except Exception as e:
+                st.error(f"❌ Error generating image: {e}")