Update app.py

app.py

@@ -62,6 +62,13 @@ 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(
@@ -69,12 +76,8 @@ def load_models():
             "Qwen/Qwen2.5-0.5B-Instruct",
             device_map="auto",
             model_kwargs={
-                "load_in_4bit": True,
-                "quantization_config": {
-                    "load_in_4bit": True,
-                    "bnb_4bit_compute_dtype": torch.float16,
-                    "bnb_4bit_quant_type": "nf4"
-                }
+                "quantization_config": quantization_config,
+                "torch_dtype": torch.float16
             }
         )
     }