Update app.py

app.py

@@ -1,86 +1,125 @@
-import os
-import time
+# app.py
 import streamlit as st
-from PIL import Image
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from gtts import gTTS
-import tempfile
-from llama_cpp import Llama
-
-# First install required package:
-# pip install llama-cpp-python
-
-# —––––––– Page Setup —–––––––
-st.set_page_config(page_title="Magic Story Generator", layout="centered")
-st.title("📖✨ Turn Images into Children's Stories")
+import os
+import time
+import torch
+from threading import Thread
 
-# —––––––– Load Models (cached) —–––––––
-@st.cache_resource(show_spinner=False)
+# Initialize models
+@st.cache_resource
 def load_models():
-    # 1) Image captioning model
-    captioner = pipeline(
-        "image-to-text",
-        model="Salesforce/blip-image-captioning-base"
+    model_name = "Qwen/Qwen3-1.7B"
+    
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name,
+        trust_remote_code=True
     )
     
-    # 2) GGUF Story Model
-    storyteller = Llama(
-        model_path="DavidAU/L3-Grand-Story-Darkness-MOE-4X8-24.9B-e32-GGUF",
-        n_ctx=2048,
-        n_threads=4,
-        n_gpu_layers=0  # Set based on your GPU capacity
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype="auto",
+        device_map="auto",
+        trust_remote_code=True
     )
-    return captioner, storyteller
+    
+    return model, tokenizer
 
-captioner, storyteller = load_models()
+def parse_thinking_output(output_ids, tokenizer, thinking_token_id=151668):
+    try:
+        index = len(output_ids) - output_ids[::-1].index(thinking_token_id)
+    except ValueError:
+        index = 0
+        
+    thinking = tokenizer.decode(output_ids[:index], skip_special_tokens=True).strip("\n")
+    content = tokenizer.decode(output_ids[index:], skip_special_tokens=True).strip("\n")
+    return thinking, content
 
-# —––––––– Main App —–––––––
-uploaded = st.file_uploader("Upload an image:", type=["jpg", "png", "jpeg"])
-if uploaded:
-    img = Image.open(uploaded).convert("RGB")
-    st.image(img, use_column_width=True)
+def generate_response(prompt, model, tokenizer):
+    messages = [{"role": "user", "content": prompt}]
+    text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        enable_thinking=True
+    )
+    
+    streamer = TextIteratorStreamer(tokenizer)
+    inputs = tokenizer([text], return_tensors="pt").to(model.device)
+    
+    generation_kwargs = dict(
+        **inputs,
+        streamer=streamer,
+        max_new_tokens=4096,
+        temperature=0.7,
+        do_sample=True
+    )
+    
+    thread = Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+    
+    full_response = ""
+    thinking_content = ""
+    for new_text in streamer:
+        full_response += new_text
+        try:
+            current_ids = tokenizer.encode(full_response, return_tensors="pt")[0]
+            thinking, content = parse_thinking_output(current_ids, tokenizer)
+            yield thinking, content
+        except:
+            yield "", full_response
 
-    # Generate caption
-    with st.spinner("🔍 Generating caption..."):
-        cap = captioner(img)
-        caption = cap[0]['generated_text']
-    st.success(f"**Caption:** {caption}")
+def text_to_speech(text):
+    tts = gTTS(text=text, lang='en', slow=False)
+    audio_file = f"audio_{int(time.time())}.mp3"
+    tts.save(audio_file)
+    return audio_file
 
-    # Generate story
-    prompt = f"""Below is an image description. Write a children's story based on it.
+# Streamlit UI
+def main():
+    st.title("🧠 Qwen3-1.7B Thinking Mode Demo")
     
-    Image Description: {caption}
-    Story:"""
+    model, tokenizer = load_models()
     
-    with st.spinner("📝 Crafting magical story..."):
-        start = time.time()
-        output = storyteller(
-            prompt=prompt,
-            max_tokens=500,
-            temperature=0.7,
-            top_p=0.9,
-            repeat_penalty=1.1
-        )
-        gen_time = time.time() - start
-        story = output['choices'][0]['text'].strip()
-        st.text(f"⏱ Generated in {gen_time:.1f}s")
-
-    # Post-process story
-    story = story.split("###")[0].strip()  # Remove any trailing artifacts
+    with st.sidebar:
+        st.header("Settings")
+        max_length = st.slider("Max Tokens", 100, 4096, 1024)
+        temperature = st.slider("Temperature", 0.1, 1.0, 0.7)
     
-    # Display story
-    st.subheader("📚 Your Magical Story")
-    st.write(story)
-
-    # Audio conversion
-    with st.spinner("🔊 Converting to audio..."):
-        try:
-            tts = gTTS(text=story, lang="en", slow=False)
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp:
-                tts.save(tmp.name)
-                st.audio(tmp.name, format="audio/mp3")
-        except Exception as e:
-            st.warning(f"⚠️ Audio conversion failed: {str(e)}")
+    prompt = st.text_area("Enter your prompt:", 
+                         "Explain quantum computing in simple terms")
+    
+    if st.button("Generate Response"):
+        with st.spinner("Generating response..."):
+            # Setup containers
+            thinking_container = st.container(border=True)
+            response_container = st.empty()
+            audio_container = st.empty()
+            
+            full_content = ""
+            current_thinking = ""
+            
+            for thinking, content in generate_response(prompt, model, tokenizer):
+                if thinking != current_thinking:
+                    thinking_container.markdown(f"**Thinking Process:**\n{thinking}")
+                    current_thinking = thinking
+                
+                if content != full_content:
+                    response_container.markdown(f"**Final Answer:**\n{content}")
+                    full_content = content
+            
+            # Add audio version
+            audio_file = text_to_speech(full_content)
+            audio_container.audio(audio_file, format='audio/mp3')
+            
+            # Add download button
+            st.download_button(
+                label="Download Response",
+                data=full_content,
+                file_name="qwen_response.txt",
+                mime="text/plain"
+            )
 
-# Footer
-st.markdown("---\n*Made with ❤️ by your friendly story wizard*")
+if __name__ == "__main__":
+    main()