Update app.py

app.py

@@ -26,7 +26,7 @@ def load_image_model():
         logger.info("Stage 1 model loaded")
         return processor, model
     except Exception as e:
-        st.error("❌ Image model failed to load")
+        st.error("❌ Failed to load image model")
         raise
 
 def stage1_generate_caption(uploaded_file):
@@ -34,7 +34,7 @@ def stage1_generate_caption(uploaded_file):
     processor, model = load_image_model()
     try:
         img = Image.open(uploaded_file).convert("RGB")
-        img.thumbnail((512, 512))
+        img.thumbnail((512, 512))  # Optimize image size
         inputs = processor(images=img, return_tensors="pt", padding=True)
         outputs = model.generate(**inputs, max_length=30)
         return processor.decode(outputs[0], skip_special_tokens=True)
@@ -47,14 +47,14 @@ def stage1_generate_caption(uploaded_file):
 # ======================
 @st.cache_resource
 def load_story_model():
-    """Load story generation model"""
+    """Load reliable story generation model"""
     try:
-        tokenizer = AutoTokenizer.from_pretrained("pranavpsv/gpt-genre-story-generator")
-        model = AutoModelForCausalLM.from_pretrained("pranavpsv/gpt-genre-story-generator", use_auth_token=True)
+        tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
+        model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
         logger.info("Stage 2 model loaded")
         return tokenizer, model
     except Exception as e:
-        st.error("❌ Story model failed to load")
+        st.error("❌ Failed to load story model")
         raise
 
 def stage2_generate_story(keyword):
@@ -62,27 +62,28 @@ def stage2_generate_story(keyword):
     tokenizer, model = load_story_model()
     
     # Optimized prompt template
-    prompt = f"""Generate a children's story with:
+    prompt = f"""Write a children's story with:
     - Theme: {keyword}
     - Characters: Animals
-    - Word count: 100 words
+    - Length: 100 words
     
-    Story: Once upon a time, a little bear named Honey discovered"""
+    Story: Once upon a time, a little bear named Honey found"""
     
     try:
         inputs = tokenizer(prompt, return_tensors="pt", max_length=100, truncation=True)
         outputs = model.generate(
             inputs.input_ids,
             max_length=300,
-            temperature=0.9,
+            temperature=0.85,
             top_k=50,
-            repetition_penalty=1.2
+            repetition_penalty=1.2,
+            pad_token_id=tokenizer.eos_token_id
         )
         full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return full_text.replace(prompt, "").strip()
     except Exception as e:
         st.error(f"Story generation failed: {str(e)}")
-        return "The animals had a wonderful adventure!"
+        return "The animals had a wonderful day playing together!"
 
 # ======================
 # Stage 3: Text-to-Speech