Update app.py

app.py

@@ -13,8 +13,8 @@ from contextlib import asynccontextmanager
 # Configuration
 MODEL_ID = "google/gemma-1.1-2b-it"
 HF_TOKEN = os.getenv("HF_TOKEN", "")
-MAX_TOKENS = 200
-DEVICE = "cpu"  # Force CPU to avoid device_map issues
+MAX_TOKENS = 150
+DEVICE = "cpu"
 PORT = int(os.getenv("PORT", 7860))
 
 # Setup logging
@@ -32,36 +32,54 @@ class ScriptGenerator:
         self.tokenizer = None
         self.model = None
         self.loaded = False
+        self.load_error = None
         
     def load_model(self):
         if self.loaded:
-            return
-        
+            return True
+            
         logger.info("Loading model...")
         try:
-            self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
+            # Load tokenizer first
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                MODEL_ID, 
+                token=HF_TOKEN
+            )
+            logger.info("✅ Tokenizer loaded")
+            
+            # Load model with simple configuration
             self.model = AutoModelForCausalLM.from_pretrained(
                 MODEL_ID,
                 torch_dtype=torch.float32,
                 token=HF_TOKEN,
-                low_cpu_mem_usage=True
+                device_map=None  # Explicitly set to None
             )
-            # Simple device assignment without device_map
+            
+            # Move to device
             self.model = self.model.to(DEVICE)
+            self.model.eval()  # Set to evaluation mode
+            
             self.loaded = True
             logger.info("✅ Model loaded successfully")
+            return True
+            
         except Exception as e:
-            logger.error(f"❌ Model loading failed: {str(e)}")
-            raise
+            self.load_error = str(e)
+            logger.error(f"❌ Model loading failed: {str(e)}", exc_info=True)
+            return False
+
+# Global generator instance
+generator = ScriptGenerator()
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    generator = ScriptGenerator()
-    generator.load_model()
+    # Load model during startup
+    success = generator.load_model()
+    if not success:
+        logger.critical("❌ Failed to load model during startup!")
     yield
 
 app = FastAPI(lifespan=lifespan)
-generator = ScriptGenerator()
 
 def extract_topic(topic_input: Union[str, List[str]]) -> str:
     """Extract topic from string or array input"""
@@ -74,40 +92,52 @@ def extract_topic(topic_input: Union[str, List[str]]) -> str:
 def generate_script(topic: str) -> str:
     """Generate script with error handling"""
     try:
+        # Check if model is loaded
+        if not generator.loaded:
+            if not generator.load_model():
+                raise Exception(f"Model failed to load: {generator.load_error}")
+        
         clean_topic = topic.strip().strip("['").strip("']").strip('"').strip("'")
         logger.info(f"🎯 Generating script for: '{clean_topic}'")
             
         prompt = (
-            f"Create a short 1-minute video script about: {clean_topic[:80]}\n\n"
-            "Structure:\n"
-            "1) Hook (5-10 seconds)\n"
-            "2) Main Content (40 seconds)\n"
-            "3) CTA (5-10 seconds)\n\n"
+            f"Create a 60-second video script about: {clean_topic[:50]}\n\n"
+            "1) Hook (10s)\n"
+            "2) Content (40s)\n" 
+            "3) CTA (10s)\n\n"
             "Script:"
         )
         
+        # Tokenize input
         inputs = generator.tokenizer(
             prompt, 
             return_tensors="pt", 
-            padding=True, 
             truncation=True,
-            max_length=512
-        ).to(DEVICE)
+            max_length=256
+        )
+        
+        # Move to device
+        inputs = {k: v.to(DEVICE) for k, v in inputs.items()}
         
-        # Generate with safer parameters
+        # Generate text
         with torch.no_grad():
             outputs = generator.model.generate(
                 **inputs,
                 max_new_tokens=MAX_TOKENS,
                 do_sample=True,
-                top_p=0.8,
+                top_p=0.9,
                 temperature=0.7,
-                pad_token_id=generator.tokenizer.eos_token_id
+                pad_token_id=generator.tokenizer.eos_token_id,
+                num_return_sequences=1
             )
         
+        # Decode output
         script = generator.tokenizer.decode(outputs[0], skip_special_tokens=True)
         clean_script = script.replace(prompt, "").strip()
         
+        if not clean_script:
+            clean_script = "Script generation completed but returned empty content."
+            
         logger.info(f"📝 Generated {len(clean_script)} characters")
         return clean_script
         
@@ -126,7 +156,6 @@ async def process_job(job_id: str, topic_input: Union[str, List[str]], callback_
             "status": "complete",
             "result": script,
             "topic": topic,
-            "original_input": topic_input,
             "script_length": len(script)
         }
         
@@ -135,35 +164,28 @@ async def process_job(job_id: str, topic_input: Union[str, List[str]], callback_
         if callback_url:
             try:
                 async with httpx.AsyncClient(timeout=30.0) as client:
-                    webhook_data = {
-                        "job_id": job_id,
-                        "status": "complete",
-                        "result": script,
-                        "topic": topic,
-                        "original_input": topic_input
-                    }
-                    
                     response = await client.post(
                         callback_url,
-                        json=webhook_data,
+                        json={
+                            "job_id": job_id,
+                            "status": "complete",
+                            "result": script,
+                            "topic": topic
+                        },
                         headers={"Content-Type": "application/json"}
                     )
-                    
                     logger.info(f"📨 Webhook status: {response.status_code}")
-                        
             except Exception as e:
                 logger.error(f"❌ Webhook failed: {str(e)}")
         
     except Exception as e:
         error_msg = f"Job failed: {str(e)}"
-        logger.error(f"❌ Job {job_id} failed: {error_msg}", exc_info=True)
+        logger.error(f"❌ Job {job_id} failed: {error_msg}")
         
         jobs[job_id] = {
             "status": "failed",
             "error": error_msg,
-            "topic": extract_topic(topic_input) if topic_input else "unknown",
-            "original_input": topic_input,
-            "script_length": 0
+            "topic": extract_topic(topic_input) if topic_input else "unknown"
         }
 
 @app.post("/api/submit")
@@ -184,10 +206,8 @@ async def submit_job(request: Request, background_tasks: BackgroundTasks):
         
         jobs[job_id] = {
             "status": "processing",
-            "result": None,
             "callback_url": callback_url,
-            "topic": topic,
-            "original_input": topic_input
+            "topic": topic
         }
         
         background_tasks.add_task(
@@ -200,12 +220,11 @@ async def submit_job(request: Request, background_tasks: BackgroundTasks):
         return JSONResponse({
             "job_id": job_id,
             "status": "queued",
-            "received_topic": topic,
-            "callback_url": callback_url
+            "topic": topic
         })
         
     except Exception as e:
-        logger.error(f"❌ Submission error: {str(e)}", exc_info=True)
+        logger.error(f"❌ Submission error: {str(e)}")
         raise HTTPException(status_code=400, detail=str(e))
 
 @app.get("/api/status/{job_id}")
@@ -235,8 +254,9 @@ async def debug_jobs():
 async def health_check():
     """Health check endpoint"""
     return {
-        "status": "healthy",
+        "status": "healthy" if generator.loaded else "unhealthy",
         "model_loaded": generator.loaded,
+        "model_error": generator.load_error,
         "total_jobs": len(jobs)
     }
 
@@ -244,11 +264,44 @@ async def health_check():
 async def test_generation():
     """Test script generation"""
     try:
-        test_topic = "healthy lifestyle tips"
+        # Check if model is loaded first
+        if not generator.loaded:
+            if not generator.load_model():
+                return {
+                    "status": "error", 
+                    "error": f"Model failed to load: {generator.load_error}"
+                }
+        
+        test_topic = "healthy lifestyle"
+        logger.info(f"🧪 Testing generation with: {test_topic}")
+        
         script = generate_script(test_topic)
-        return {"status": "success", "topic": test_topic, "script": script}
+        
+        return {
+            "status": "success", 
+            "topic": test_topic, 
+            "script_length": len(script),
+            "script_preview": script[:200] + "..." if len(script) > 200 else script
+        }
+        
     except Exception as e:
-        return {"status": "error", "error": str(e)}
+        logger.error(f"❌ Test generation failed: {str(e)}", exc_info=True)
+        return {
+            "status": "error", 
+            "error": str(e),
+            "model_loaded": generator.loaded,
+            "model_error": generator.load_error
+        }
+
+@app.get("/test/model")
+async def test_model():
+    """Test if model loads correctly"""
+    return {
+        "model_loaded": generator.loaded,
+        "model_error": generator.load_error,
+        "has_tokenizer": generator.tokenizer is not None,
+        "has_model": generator.model is not None
+    }
 
 if __name__ == "__main__":
     uvicorn.run(