Update app.py

app.py

@@ -6,54 +6,26 @@ import gradio as gr
 # Configuration
 MODEL_ID = "google/gemma-1.1-7b-it"
 HF_TOKEN = os.getenv("HF_TOKEN")
-MAX_TOKENS = 200  # Reduced for stability
+MAX_TOKENS = 150  # Reduced for CPU
 
-# Initialize components
-tokenizer = None
-model = None
-
-def load_model():
-    """Lazy-load the model to avoid immediate memory issues"""
-    global tokenizer, model
-    if model is None:
-        print("🚀 Loading model...")
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID,
-            device_map="auto",
-            torch_dtype=torch.float16,
-            token=HF_TOKEN
-        )
-        print("✅ Model loaded!")
+# Load model (CPU-only)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    device_map="cpu",
+    torch_dtype=torch.float32,
+    token=HF_TOKEN
+)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
 
-def generate_script(topic):
-    """The prediction function that handles requests"""
-    try:
-        load_model()
-        prompt = f"Generate a short YouTube script about {topic}:"
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=MAX_TOKENS,
-            temperature=0.7
-        )
-        return tokenizer.decode(outputs[0], skip_special_tokens=True)
-    except Exception as e:
-        print(f"⚠️ Error: {str(e)}")
-        return f"Generation failed: {str(e)}"
+def predict(topic):
+    prompt = f"Create a short script about {topic}:"
+    inputs = tokenizer(prompt, return_tensors="pt")
+    outputs = model.generate(**inputs, max_new_tokens=MAX_TOKENS)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Create Gradio app
-app = gr.Interface(
-    fn=generate_script,
+gr.Interface(
+    fn=predict,
     inputs=gr.Textbox(label="Topic"),
     outputs=gr.Textbox(label="Script"),
-    title="Gemma-7B Script Generator"
-)
-
-# Launch with explicit API configuration
-app.launch(
-    server_name="0.0.0.0",
-    server_port=7860,
-    enable_api=True,  # CRITICAL FOR API
-    share=False
-)
+    api_name="predict"
+).launch(server_name="0.0.0.0")