Update app.py

app.py

@@ -1,48 +1,115 @@
 import streamlit as st
 import torch
-from transformers import pipeline
+from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 from diffusers import StableDiffusionPipeline
-from PIL import Image
+from PIL import Image, ImageDraw, ImageFont
 
-# Set device (CPU or GPU)
+# Check for GPU availability
 device = "cuda" if torch.cuda.is_available() else "cpu"
+st.write(f"Using device: {device}")  # Debug message
 
-# Load text generation model
+# Load text model (TinyLlama) with caching
 @st.cache_resource
 def load_text_model():
-    return pipeline("text-generation", model="gpt2")
+    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
 
-text_generator = load_text_model()
+story_generator = load_text_model()
 
-# Load image generation model
+# Load image model (Stable Diffusion) with caching
 @st.cache_resource
 def load_image_model():
-    pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base")
-    pipe.to(device)
-    pipe.to(torch.float16)  # Use float16 for speed
-    if device == "cuda":
-        pipe = torch.compile(pipe)  # Optimize for GPU
-    return pipe
+    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:"
+    
+    try:
+        st.write("⏳ Generating story...")
+        story_output = story_generator(
+            formatted_prompt,
+            max_length=120,  # Reduced 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)
+
+    try:
+        font = ImageFont.truetype("arial.ttf", 20)
+    except IOError:
+        font = ImageFont.load_default()
+
+    text_bbox = draw.textbbox((0, 0), text, font=font)
+    text_width = text_bbox[2] - text_bbox[0]
+    text_height = text_bbox[3] - text_bbox[1]
+
+    bubble_width, bubble_height = text_width + 30, text_height + 20
+    bubble_x, bubble_y = position
+
+    draw.ellipse([bubble_x, bubble_y, bubble_x + bubble_width, bubble_y + bubble_height], fill="white", outline="black")
+    draw.text((bubble_x + 15, bubble_y + 10), text, font=font, fill="black")
+
+    return image
+
 # Streamlit UI
-st.title("🤖 AI Comic Story Generator")
+st.title("🦸‍♂️ AI Comic Story Generator")
 st.write("Enter a prompt to generate a comic-style story and image!")
 
-# Input field for story prompt
-story_prompt = st.text_input("📝 Enter your story prompt:", "")
+# User input
+user_prompt = st.text_input("📝 Enter your story prompt:")
+
+if user_prompt:
+    st.subheader("📖 AI-Generated Story")
+    generated_story = generate_story(user_prompt)
+    st.write(generated_story)
 
-if story_prompt:
-    with st.spinner("⏳ Generating story..."):
-        story = text_generator(story_prompt, max_length=100, num_return_sequences=1)[0]["generated_text"]
-    st.success("✅ Story generated successfully!")
-    st.write(story)
+    st.subheader("🖼️ AI-Generated Image")
+    
+    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=10).images[0]  # Reduced steps from 12 → 10
+                image = image.resize((512, 512))  # Resize to smaller 512x512
+                st.write("✅ Image generated successfully!")
 
-    # Generate image
-    with st.spinner("🎨 Generating image..."):
-        image = image_generator(story_prompt, num_inference_steps=15).images[0]  # Reduced from 30 → 15
-        image = image.resize((512, 512))  # Resize to 512x512 to make it smaller
-        st.image(image, caption="🖼��� AI-Generated Comic Image", use_column_width=False)
+                # Extract first sentence (max 50 characters) for speech bubble
+                speech_text = generated_story.split(".")[0][:50]
+                image_with_bubble = add_speech_bubble(image, speech_text, position=(50, 50))
 
-st.write("🚀 Optimized for speed & performance!")
+                st.image(image_with_bubble, caption="Generated Comic Image", use_container_width=True)
+            except Exception as e:
+                st.error(f"❌ Error generating image: {e}")