Update app.py

app.py

@@ -12,25 +12,20 @@ def generate_caption(image_file):
     caption = caption_results[0]['generated_text']
     return caption
 
-
 def generate_story(caption):
-
-    story_generator = pipeline("text-generation", model="openai-community/gpt2")
-    prompt = f"Please based on following image caption: '{caption}', generate a complete fairy tale story for children with at least 100 words. "
-    result = story_generator(prompt, max_length=300, num_return_sequences=1)
+    story_generator = pipeline("text-generation", model="deepseek-ai/DeepSeek-R1-Distill-Qwen-32B")
+    messages = [
+        {
+            "role": "user",
+            "content": f"Please based on following image caption: '{caption}', generate a complete fairy tale story for children with at least 100 words and max 300 words"
+        }
+    ]
+    result = story_generator(messages, max_length=300, num_return_sequences=1)
     story = result[0]['generated_text']
-    
-    if len(story.split()) < 100:
-        additional = story_generator(prompt, max_length=350, num_return_sequences=1)[0]['generated_text']
-        story += " " + additional
     return story
 
-# ----------------------------
-# generate_illustration
-# ----------------------------
 @st.cache_resource
 # def load_image_generator():
-
 #     device = "cuda" if torch.cuda.is_available() else "cpu"
 #     torch_dtype = torch.float16 if device == "cuda" else torch.float32
 #     pipe = DiffusionPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5")
@@ -43,7 +38,6 @@ def generate_story(caption):
 #     generated_image = image_result.images[0]
 #     return generated_image
 
-
 def text_to_speech(text, output_file="output.mp3"):
     tts = gTTS(text=text, lang="en")
     tts.save(output_file)
@@ -70,7 +64,6 @@ def main():
         
         # with st.spinner("Generating illustration..."):
         #     illustration = generate_illustration(story[:200])
-        
         # st.write("### Story Illustrations：")
         # st.image(illustration, caption="Story Illustrations", use_container_width=True)