Update app.py

app.py

@@ -1,34 +1,23 @@
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import gradio as gr
 
-# Load model (cache on first run)
 model_id = "google/gemma-1.1-7b-it"
+
+# CPU-specific config
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    device_map="auto",
-    torch_dtype=torch.bfloat16  # Optimized for NVIDIA
-)
-
-def generate_script(topic):
-    prompt = f"""Generate a viral YouTube Short script about {topic} with:
-    1) HOOK: Controversial opening (5 words max)
-    2) BODY: 3 scientific facts
-    3) CTA: Actionable challenge
-    
-    Script:"""
-    
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
-    outputs = model.generate(**inputs, max_new_tokens=300)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# Gradio UI
-demo = gr.Interface(
-    fn=generate_script,
-    inputs=gr.Textbox(label="Topic"),
-    outputs=gr.Textbox(label="Generated Script"),
-    title="Gemma-7B Script Generator"
+    device_map="cpu",
+    torch_dtype=torch.float32,  # Required for CPU
+    load_in_8bit=True  # Reduces RAM usage by 2x
 )
 
-demo.launch(server_port=7860, share=True)
+def generate(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=150,  # Must stay under 200
+        do_sample=True,
+        temperature=0.7
+    )
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)