Update app.py

app.py

@@ -2,52 +2,111 @@ import os
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
+from typing import Dict, Any
 
 # Configuration
-MODEL_ID = "google/gemma-1.1-2b-it"  # Using smaller 2B version
-HF_TOKEN = os.getenv("HF_TOKEN")
-MAX_TOKENS = 80  # Conservative limit
-
-def load_model():
-    """Simplified model loading that works in Spaces"""
-    print("🔄 Loading model...")
-    
-    # Load tokenizer
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
-    
-    # Explicit CPU-only loading
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        torch_dtype=torch.float32,  # Required for CPU
-        token=HF_TOKEN
-    ).to('cpu')  # Explicit CPU placement
-    
-    print("✅ Model loaded successfully!")
-    return tokenizer, model
-
-tokenizer, model = load_model()
-
-def predict(topic):
-    """Memory-safe generation"""
+MODEL_ID = "google/gemma-1.1-2b-it"
+HF_TOKEN = os.getenv("HF_TOKEN", "")  # Default empty if not set
+MAX_TOKENS = 80
+MAX_INPUT_LENGTH = 100
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+class ScriptGenerator:
+    def __init__(self):
+        self.tokenizer = None
+        self.model = None
+        self.loaded = False
+        
+    def load_model(self):
+        """Safe model loading with progress tracking"""
+        if self.loaded:
+            return
+            
+        print("🔄 Loading model...")
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
+            
+            self.model = AutoModelForCausalLM.from_pretrained(
+                MODEL_ID,
+                torch_dtype=torch.float32 if DEVICE == "cpu" else torch.float16,
+                device_map="auto" if DEVICE == "cuda" else None,
+                token=HF_TOKEN
+            ).to(DEVICE)
+            
+            self.loaded = True
+            print("✅ Model loaded successfully!")
+        except Exception as e:
+            print(f"❌ Model loading failed: {str(e)}")
+            raise
+
+generator = ScriptGenerator()
+
+def predict(topic: str) -> str:
+    """Generate script with proper error handling"""
     try:
+        # Input validation
+        if not topic or len(topic) > MAX_INPUT_LENGTH:
+            return f"Topic must be 1-{MAX_INPUT_LENGTH} characters"
+            
+        if not generator.loaded:
+            generator.load_model()
+            
         prompt = f"Create a short script about {topic}:\n1) Hook\n2) Point\n3) CTA\n\nScript:"
-        inputs = tokenizer(prompt, return_tensors="pt").to('cpu')
         
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=MAX_TOKENS,
-            temperature=0.7
-        )
+        with torch.no_grad():
+            inputs = generator.tokenizer(prompt, return_tensors="pt").to(DEVICE)
+            outputs = generator.model.generate(
+                **inputs,
+                max_new_tokens=MAX_TOKENS,
+                temperature=0.7,
+                do_sample=True
+            )
+            
+        return generator.tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        return tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
+    except torch.cuda.OutOfMemoryError:
+        return "Error: GPU out of memory - try a shorter input"
     except Exception as e:
         return f"Error: {str(e)}"
 
-# Minimal interface
-gr.Interface(
+def api_predict(data: Dict[str, Any]) -> Dict[str, Any]:
+    """Dedicated API endpoint with standardized response"""
+    try:
+        topic = data.get("topic", "")
+        result = predict(topic)
+        return {
+            "success": not result.startswith("Error"),
+            "result": result,
+            "error": result if result.startswith("Error") else None
+        }
+    except Exception as e:
+        return {
+            "success": False,
+            "result": None,
+            "error": str(e)
+        }
+
+# Gradio Interface with explicit API
+interface = gr.Interface(
     fn=predict,
-    inputs=gr.Textbox(label="Topic"),
-    outputs=gr.Textbox(label="Script", lines=4),
-    api_name="predict"
-).launch(server_name="0.0.0.0")
+    inputs=gr.Textbox(label="Topic", placeholder="Enter your script topic..."),
+    outputs=gr.Textbox(label="Generated Script", lines=5),
+    title="Gemma Script Generator",
+    description="Generate marketing scripts using Gemma 2B"
+)
+
+# Mount both UI and API
+app = gr.mount_gradio_app(
+    gr.App(),
+    interface,
+    path="/"
+)
+app.add_api_route("/api/predict", api_predict, methods=["POST"])
+
+if __name__ == "__main__":
+    generator.load_model()
+    app.launch(
+        server_name="0.0.0.0",
+        server_port=int(os.getenv("PORT", 7860)),
+        share=False
+    )