Update app.py

app.py

@@ -1,24 +1,28 @@
+
 import os
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import BitsAndBytesConfig
 import gradio as gr
 
 # Configuration
 MODEL_ID = "google/gemma-1.1-7b-it"
-HF_TOKEN = os.getenv("HF_TOKEN")  # Make sure this is set in Space secrets
-MAX_TOKENS = 300
+HF_TOKEN = os.getenv("HF_TOKEN")
+MAX_TOKENS = 250  # Reduced for stability
 
-# Initialize model
-tokenizer = AutoTokenizer.from_pretrained(
-    MODEL_ID,
-    token=HF_TOKEN
+# 4-bit quantization config
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_quant_type="nf4"
 )
 
-# Load model with CPU fallback
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
     device_map="auto",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    quantization_config=quant_config,
     token=HF_TOKEN
 )
 
@@ -30,9 +34,7 @@ def generate_script(topic):
     
     Script:"""
     
-    # Use current device (automatically handles CPU/GPU)
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     
     outputs = model.generate(
         **inputs,
@@ -45,39 +47,13 @@ def generate_script(topic):
 
 # Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown("## 🚀 Gemma-7B Script Generator")
+    gr.Markdown("## 🎥 Optimized Gemma-7B Generator")
     with gr.Row():
-        topic_input = gr.Textbox(
-            label="Enter your topic",
-            placeholder="e.g., 'intermittent fasting benefits'"
-        )
+        topic_input = gr.Textbox(label="Topic", placeholder="e.g., 'cold shower benefits'")
         generate_btn = gr.Button("Generate", variant="primary")
-    
-    script_output = gr.Textbox(
-        label="Generated Script",
-        interactive=False,
-        lines=10
-    )
-    
-    # Examples
-    gr.Examples(
-        examples=[
-            ["Why cold showers boost metabolism"],
-            ["3 workout myths debunked by science"],
-            ["The truth about protein timing"]
-        ],
-        inputs=topic_input
-    )
+    script_output = gr.Textbox(label="Script", lines=8)
     
     generate_btn.click(
         fn=generate_script,
         inputs=topic_input,
-        outputs=script_output
-    )
-
-# Launch with production settings
-demo.launch(
-    server_name="0.0.0.0",
-    server_port=7860,
-    share=False
-)
+        outputs=script_output