Update app.py

app.py

@@ -2,7 +2,7 @@ import nest_asyncio
 nest_asyncio.apply()
 
 import streamlit as st
-from transformers import pipeline, BitsAndBytesConfig
+from transformers import pipeline
 import torch
 from gtts import gTTS
 import io
@@ -62,24 +62,9 @@ st.header("Turn Your Image to a Short Audio Story for Children")
 # Model loading
 @st.cache_resource
 def load_models():
-    # Configure 4-bit quantization properly
-    quantization_config = BitsAndBytesConfig(
-        load_in_4bit=True,
-        bnb_4bit_compute_dtype=torch.float16,
-        bnb_4bit_quant_type="nf4"
-    )
-    
     return {
         "img_model": pipeline("image-to-text", "cnmoro/tiny-image-captioning"),
-        "story_model": pipeline(
-            "text-generation", 
-            "Qwen/Qwen2.5-0.5B-Instruct",
-            device_map="auto",
-            model_kwargs={
-                "quantization_config": quantization_config,
-                "torch_dtype": torch.float16
-            }
-        )
+        "story_model": pipeline("text-generation", "Qwen/Qwen2.5-0.5B-Instruct")
     }
 
 models = load_models()
@@ -96,7 +81,7 @@ def text2story(text):
     ]
     response = models["story_model"](
         messages,
-        max_new_tokens=128,
+        max_new_tokens=100,
         do_sample=True,
         temperature=0.7
     )[0]["generated_text"]