Update app.py

app.py

@@ -1,109 +1,61 @@
-import os
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
-from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse
-import logging
-import uvicorn
-
-# Configure logging
-logging.basicConfig(
-    level=logging.INFO,
-    format='%(asctime)s - %(levelname)s - %(message)s'
-)
-logger = logging.getLogger(__name__)
-
-# Configuration
-MODEL_ID = "google/gemma-1.1-2b-it"
-HF_TOKEN = os.getenv("HF_TOKEN", "")
-MAX_TOKENS = 400
-DEVICE = "cpu"
-PORT = int(os.getenv("PORT", 7860))
-
-class ScriptGenerator:
-    def __init__(self):
-        self.tokenizer = None
-        self.model = None
-        self.generation_config = None
-        self.loaded = False
-        
-    def load_model(self):
-        if self.loaded: return
-        
-        logger.info("Loading model...")
-        try:
-            self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
-            self.model = AutoModelForCausalLM.from_pretrained(
-                MODEL_ID,
-                torch_dtype=torch.float32,
-                device_map=None,
-                token=HF_TOKEN,
-                low_cpu_mem_usage=True
-            ).to(DEVICE)
-            
-            self.generation_config = GenerationConfig(
-                max_new_tokens=MAX_TOKENS,
-                do_sample=True,
-                top_p=0.9,
-                num_beams=1,
-                no_repeat_ngram_size=2,
-                pad_token_id=self.tokenizer.eos_token_id
-            )
-            
-            self.loaded = True
-            logger.info("Model loaded | Port: %s", PORT)
-        except Exception as e:
-            logger.error("Load failed: %s", str(e))
-            raise
-
-generator = ScriptGenerator()
-
-app = FastAPI()
-
-@app.on_event("startup")
-def startup():
-    generator.load_model()
+from fastapi import BackgroundTasks
+import httpx
+import uuid
+
+jobs = {}  # Stores ongoing jobs
+
+@app.post("/api/submit")
+async def submit_job(
+    request: Request, 
+    background_tasks: BackgroundTasks
+):
+    data = await request.json()
+    job_id = str(uuid.uuid4())
+    
+    # Store job details
+    jobs[job_id] = {
+        "status": "processing",
+        "result": None,
+        "callback_url": data.get("callback_url")  # n8n webhook URL
+    }
+    
+    # Start background task
+    background_tasks.add_task(
+        process_job, 
+        job_id, 
+        data["topic"]
+    )
+    
+    return {"job_id": job_id, "status": "queued"}
 
-@app.post("/api/predict")
-async def predict(request: Request):
+async def process_job(job_id: str, topic: str):
     try:
-        data = await request.json()
-        topic = data.get("topic", "")
-        
-        if isinstance(topic, list):
-            topic = topic[0] if len(topic) > 0 else ""
-        topic = str(topic).strip()
-        
-        logger.info("Processing: %.30s...", topic)
-        
-        inputs = generator.tokenizer(
-            f"Create 1-minute script about {topic}:\n1) Hook\n2) Main\n3) CTA\n\nScript:",
-            return_tensors="pt"
-        ).to(DEVICE)
-        
-        outputs = generator.model.generate(
-            **inputs,
-            generation_config=generator.generation_config
-        )
-        
-        script = generator.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return JSONResponse({"result": script})
-        
+        script = generate_script(topic)  # Your existing function
+        jobs[job_id]["status"] = "complete"
+        jobs[job_id]["result"] = script
+        
+        # Send back to n8n via webhook
+        if jobs[job_id]["callback_url"]:
+            async with httpx.AsyncClient() as client:
+                await client.post(
+                    jobs[job_id]["callback_url"],
+                    json={
+                        "job_id": job_id,
+                        "status": "complete",
+                        "result": script
+                    },
+                    timeout=30.0
+                )
+                
     except Exception as e:
-        logger.error("API error: %s", str(e))
-        return JSONResponse({"error": str(e)}, status_code=500)
-
-if __name__ == "__main__":
-    # Hugging Face Spaces compatibility
-    if os.getenv("SPACES", "false").lower() == "true":
-        os.environ["GRADIO_SERVER_PORT"] = str(PORT)
-        os.environ["GRADIO_SERVER_NAME"] = "0.0.0.0"
-    
-    uvicorn.run(
-        app,
-        host="0.0.0.0",
-        port=PORT,
-        log_level="info",
-        workers=1,
-        timeout_keep_alive=30
-    )
+        jobs[job_id]["status"] = "failed"
+        jobs[job_id]["error"] = str(e)
+        if jobs[job_id]["callback_url"]:
+            await httpx.post(
+                jobs[job_id]["callback_url"],
+                json={
+                    "job_id": job_id,
+                    "status": "failed",
+                    "error": str(e)
+                }
+            )