Update app.py

app.py

@@ -5,41 +5,60 @@ from PIL import Image
 from io import BytesIO
 from huggingface_hub import InferenceApi, login
 from transformers import pipeline
+import torch
+from transformers import Qwen2_5OmniForConditionalGeneration, Qwen2_5OmniTokenizer
 from gtts import gTTS
 import tempfile
 
 # —––––––– Page Config —–––––––
-st.set_page_config(page_title="Magic Story Generator", layout="centered")
-st.title("📖✨ Turn Images into Children's Stories")
+st.set_page_config(page_title="Magic Story Generator (Qwen2.5)", layout="centered")
+st.title("📖✨ Turn Images into Children's Stories (Qwen2.5-Omni-7B)")
 
 # —––––––– Load Clients & Pipelines (cached) —–––––––
 @st.cache_resource(show_spinner=False)
 def load_clients():
     hf_token = st.secrets["HF_TOKEN"]
-    # authenticate so transformers can pick up your token
+    # Authenticate for HF Hub
     os.environ["HUGGINGFACEHUB_API_TOKEN"] = hf_token
     login(hf_token)
 
-    # BLIP captioning via Hugging Face Inference API
+    # 1) BLIP captioning via HF Inference API
     caption_client = InferenceApi(
         repo_id="Salesforce/blip-image-captioning-base",
         token=hf_token
     )
 
-    # Instruction-tuned story generator: Flan-T5
+    # 2) Qwen2.5-Omni story generator
     t0 = time.time()
+    model = Qwen2_5OmniForConditionalGeneration.from_pretrained(
+        "Qwen/Qwen2.5-Omni-7B",
+        device_map="auto",
+        torch_dtype=torch.bfloat16,
+        attn_implementation="flash_attention_2",
+        trust_remote_code=True
+    )
+    tokenizer = Qwen2_5OmniTokenizer.from_pretrained(
+        "Qwen/Qwen2.5-Omni-7B",
+        trust_remote_code=True
+    )
     storyteller = pipeline(
         task="text2text-generation",
-        model="google/flan-t5-small",
-        device=-1,               # CPU
-        max_length=150           # prompt + generation cap
+        model=model,
+        tokenizer=tokenizer,
+        device_map="auto",
+        temperature=0.7,
+        top_p=0.9,
+        repetition_penalty=1.2,
+        no_repeat_ngram_size=3,
+        max_new_tokens=120
     )
-    st.text(f"✅ Story model loaded in {time.time() - t0:.1f}s")
+    load_time = time.time() - t0
+    st.text(f"✅ Story model loaded in {load_time:.1f}s (cached thereafter)")
+
     return caption_client, storyteller
 
 caption_client, storyteller = load_clients()
 
-
 # —––––––– Helpers —–––––––
 def generate_caption(img: Image.Image) -> str:
     buf = BytesIO()
@@ -49,32 +68,33 @@ def generate_caption(img: Image.Image) -> str:
         return resp[0].get("generated_text", "").strip()
     return ""
 
+
 def generate_story(caption: str) -> str:
     prompt = (
-        "You are a creative children’s-story author.\n"
+        "You are a creative children's-story author.\n"
         f"Image description: “{caption}”\n\n"
         "Write a coherent 50–100 word story\n"
     )
     t0 = time.time()
-    out = storyteller(prompt, max_new_tokens=120, temperature=0.7, top_p=0.9)[0]["generated_text"]
-    st.text(f"⏱ Generated in {time.time() - t0:.1f}s")
-    story = out.strip()
+    outputs = storyteller(prompt)
+    gen_time = time.time() - t0
+    st.text(f"⏱ Generated in {gen_time:.1f}s on GPU/CPU")
 
-    # Truncate to at most 100 words
+    story = outputs[0]["generated_text"].strip()
+    # Enforce ≤100 words
     words = story.split()
     if len(words) > 100:
         story = " ".join(words[:100])
-        if not story.endswith("."):
-            story += "."
+        if not story.endswith('.'):
+            story += '.'
     return story
 
-
 # —––––––– Main App —–––––––
 uploaded = st.file_uploader("Upload an image:", type=["jpg","png","jpeg"])
 if uploaded:
     img = Image.open(uploaded).convert("RGB")
     if max(img.size) > 2048:
-        img.thumbnail((2048, 2048))
+        img.thumbnail((2048,2048))
     st.image(img, use_container_width=True)
 
     with st.spinner("🔍 Generating caption..."):