Update app.py

app.py

@@ -1,114 +1,48 @@
 import streamlit as st
 import torch
-from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
+from transformers import pipeline
 from diffusers import StableDiffusionPipeline
-from PIL import Image, ImageDraw, ImageFont
+from PIL import Image
 
-# Check for GPU availability
+# Set device (CPU or GPU)
 device = "cuda" if torch.cuda.is_available() else "cpu"
-st.write(f"Using device: {device}")  # Debug message
 
-# Load text model (TinyLlama) with error handling
+# Load text generation model
 @st.cache_resource
 def load_text_model():
-    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
+    return pipeline("text-generation", model="gpt2")
 
-story_generator = load_text_model()
+text_generator = load_text_model()
 
-# Load image model (Stable Diffusion) with error handling
+# Load image generation model
 @st.cache_resource
 def load_image_model():
-    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
+    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
 
 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=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)
-
-    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!")
 
-# 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)
+# Input field for story prompt
+story_prompt = st.text_input("📝 Enter your story prompt:", "")
 
-    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=30).images[0]
-                st.write("✅ Image generated successfully!")
+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)
 
-                # 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))
+    # 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)
 
-                st.image(image_with_bubble, caption="Generated Comic Image", use_container_width=True)
-            except Exception as e:
-                st.error(f"❌ Error generating image: {e}")
+st.write("🚀 Optimized for speed & performance!")