Update app.py

app.py

@@ -27,18 +27,17 @@ def load_pipelines():
         model="Salesforce/blip-image-captioning-base",
         device=-1  # CPU; set to 0+ for GPU
     )
-    # 2) Story-generation pipeline (DeepSeek-R1-Distill-Qwen)
+    # 2) Story-generation pipeline (StoryTeller)
     storyteller = pipeline(
         task="text2text-generation",
-        model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
-        tokenizer="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
-        trust_remote_code=True,
+        model="pranavpsv/StoryTeller",
+        tokenizer="pranavpsv/StoryTeller",
         device=-1,                # CPU; set to 0+ for GPU
-        temperature=0.6,
+        temperature=0.7,
         top_p=0.9,
-        repetition_penalty=1.1,
-        no_repeat_ngram_size=2,
-        max_new_tokens=120
+        repetition_penalty=1.2,
+        no_repeat_ngram_size=3,
+        max_new_tokens=150
     )
     return captioner, storyteller
 
@@ -49,7 +48,7 @@ uploaded = st.file_uploader("Upload an image:", type=["jpg", "png", "jpeg"])
 if uploaded:
     # Load and display the image
     img = Image.open(uploaded).convert("RGB")
-    st.image(img, use_container_width=True)
+    st.image(img, use_column_width=True)
 
     # Generate caption
     with st.spinner("🔍 Generating caption..."):
@@ -94,4 +93,3 @@ if uploaded:
 # Footer
 st.markdown("---\n*Made with ❤️ by your friendly story wizard*")
 
-