Update app.py

app.py

@@ -15,11 +15,10 @@ except ImportError:
     import torch
 
 # Load the image captioning model
-caption_model = pipeline("image-to-text", model="kha-white/manga-ocr-base")
+caption_model = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base")
 
-# Load the text generation model
-text_generation_model = AutoModelForCausalLM.from_pretrained("gpt2")
-tokenizer = AutoTokenizer.from_pretrained("gpt2")
+# Load the DeepSeek model for story generation
+story_generator = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1", trust_remote_code=True)
 
 def generate_caption(image):
     # Generate the caption for the uploaded image
@@ -27,11 +26,10 @@ def generate_caption(image):
     return caption
 
 def generate_story(caption):
-    # Generate the story based on the caption
+    # Generate the story based on the caption using the DeepSeek model
     prompt = f"Imagine you are a storyteller for young children. Based on the image described as '{caption}', create a short and interesting story for children aged 3-10. Keep it positive and happy in tone."
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    output = text_generation_model.generate(input_ids, max_length=200, num_return_sequences=1)
-    story = tokenizer.decode(output[0], skip_special_tokens=True)
+    messages = [{"role": "user", "content": prompt}]
+    story = story_generator(messages)[0]["generated_text"]
     return story
 
 def convert_to_audio(story):
@@ -60,7 +58,7 @@ def main():
         st.subheader("Generated Caption:")
         st.write(caption)
         
-        # Generate the story based on the caption
+        # Generate the story based on the caption using the DeepSeek model
         story = generate_story(caption)
         st.subheader("Generated Story:")
         st.write(story)