Update comfy_engine.py

comfy_engine.py

@@ -68,78 +68,141 @@ def complete_task(task_id, video_url=None, error=None):
     save_history(data)
 
 # ==========================================
-# ⚙️ CORE ENGINE LOGIC
+# 🗄️ SMART HISTORY MANAGER (Live Tracking)
+# ==========================================
+def load_history():
+    if os.path.exists("history.json"):
+        try:
+            with open("history.json", "r", encoding="utf-8") as f:
+                return json.load(f)
+        except:
+            pass
+    return []
+
+def save_history(data):
+    with open("history.json", "w", encoding="utf-8") as f:
+        json.dump(data[:50], f, indent=4) # 50 वीडियोस की लिमिट
+
+def init_task(task_id, prompt):
+    data = load_history()
+    # 👈 फिक्स 1: शुरुआत में इसे 'queued' रखो ताकि UI पर पीला वाला IN QUEUE दिखे
+    new_entry = {
+        "id": task_id,
+        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        "prompt": prompt,
+        "status": "queued", 
+        "progress": 0,
+        "url": None
+    }
+    data.insert(0, new_entry)
+    save_history(data)
+
+# 👈 फिक्स 2: प्रोग्रेस के साथ-साथ स्टेटस भी अपडेट करने का ऑप्शन
+def update_task_progress(task_id, progress, status=None):
+    data = load_history()
+    for item in data:
+        if item["id"] == task_id:
+            item["progress"] = progress
+            if status:
+                item["status"] = status
+            break
+    save_history(data)
+
+def complete_task(task_id, video_url=None, error=None):
+    data = load_history()
+    for item in data:
+        if item["id"] == task_id:
+            if error:
+                item["status"] = "failed"
+                item["error"] = error
+            else:
+                item["status"] = "done"
+                item["progress"] = 100
+                item["url"] = video_url
+            break
+    save_history(data)
+
+
+# ==========================================
+# 🚀 THE BULLETPROOF BACKGROUND WATCHER
 # ==========================================
-def inject_params(req: dict) -> dict:
-    p = json.loads(json.dumps(WORKFLOW_TEMPLATE))
-    p["89"]["inputs"]["text"] = req.get("prompt", "")
-    ratio = req.get("aspect_ratio", "16:9")
-    quality = req.get("quality", "480p")
-    width, height = RESOLUTIONS.get(quality, {}).get(ratio, (848, 480))
-    p["74"]["inputs"]["width"] = width
-    p["74"]["inputs"]["height"] = height
-    p["74"]["inputs"]["length"] = req.get("frames", 81)
-    p["88"]["inputs"]["fps"] = req.get("fps", 16)
-    
-    # ❌ यहाँ से मैंने steps और cfg वाली दोनों लाइनें डिलीट कर दी हैं!
-    # अब यह कोड सीधा workflow.json की मास्टर सेटिंग्स उठाएगा।
-    return p
-
-def extract_video_url(history: dict, token: str) -> str:
-    outputs = history.get("outputs", {})
-    for _, node_out in outputs.items():
-        for key in ("videos", "images", "files"):
-            if key in node_out and node_out[key]:
-                it = node_out[key][0]
-                q = urlencode(it)
-                return f"/api/video?{q}&token={token}"
-    raise RuntimeError("Video not found")
-
-async def queue_prompt(req: dict):
-    token = str(uuid.uuid4())
-    client_id = str(uuid.uuid4())
-    prompt_data = inject_params(req)
-    
-    async with httpx.AsyncClient() as client:
-        resp = await client.post(f"http://{COMFY_HOST}/prompt?token={token}", 
-                               json={"prompt": prompt_data, "client_id": client_id}, 
-                               timeout=30.0)
-        prompt_id = resp.json().get("prompt_id")
-        if not prompt_id: raise Exception("Failed to queue on AMD Server")
-        
-    return prompt_id, client_id, token, len(prompt_data)
-
-# 🚀 THE BACKGROUND WATCHER (यह फोन बंद होने पर भी चलता रहेगा)
 async def background_watcher(task_id, client_id, token, total_nodes):
     seen = set()
     start_t = time.time()
     progress_fake = 0
     ws_url = f"ws://{COMFY_HOST}/ws?clientId={client_id}&token={token}"
-    
+
+    # 1. WEBSOCKET LOOP (Live Progress के लिए)
     try:
-        async with websockets.connect(ws_url, ping_interval=20) as ws:
+        # 👈 टाइमआउट बढ़ा दिया है ताकि जल्दी ना कटे
+        async with websockets.connect(ws_url, ping_interval=30, ping_timeout=120) as ws:
             while True:
                 msg_raw = await ws.recv()
+                
                 if isinstance(msg_raw, (bytes, bytearray)):
                     if progress_fake < 95 and (time.time() - start_t) > 2:
                         progress_fake = min(95, progress_fake + 1)
-                        update_task_progress(task_id, progress_fake) # लाइव प्रोग्रेस सेव
+                        update_task_progress(task_id, progress_fake, "processing")
                     continue
-                
+                    
                 msg = json.loads(msg_raw)
-                if msg.get("type") == "executing":
+                
+                # जब GPU असली रेंडरिंग शुरू करे
+                if msg.get("type") == "execution_start":
+                    update_task_progress(task_id, progress_fake, "processing")
+                    
+                elif msg.get("type") == "executing":
                     node = msg.get("data", {}).get("node")
-                    if node is None: break # रेंडरिंग ख़तम
+                    if node is None:
+                        break # रेंडरिंग ख़तम
+                        
                     if node not in seen:
                         seen.add(node)
-                        p_real = int((len(seen) / total_nodes) * 100)
-                        progress_fake = max(progress_fake, p_real)
-                        update_task_progress(task_id, progress_fake) # लाइव प्रोग्रेस सेव
-    except Exception:
-        pass # अगर कनेक्शन टूटा तो भी कोई बात नहीं, हम हिस्ट्री चेक कर लेंगे!
-
-    # 🎬 फाइनल चेक और सेविंग
-    await asyncio.sleep(2)
+                        
+                    p_real = int((len(seen) / total_nodes) * 100)
+                    progress_fake = max(progress_fake, p_real)
+                    update_task_progress(task_id, progress_fake, "processing")
+                    
+    except Exception as e:
+        print(f"⚠️ WS Drop for {task_id}: {e} (Switching to HTTP Polling...)")
+        pass # अगर फोन कट भी गया, तो नीचे वाला जासूस संभाल लेगा!
+
+    # 2. HTTP POLLING FALLBACK (असली ब्रह्मास्त्र)
+    # यह तब तक बैकएंड से पूछता रहेगा जब तक वीडियो सच में बन नहीं जाता
+    while True:
+        try:
+            async with httpx.AsyncClient() as client:
+                # A. सबसे पहले चेक करो कि वीडियो बनकर हिस्ट्री में आ गया क्या?
+                h_resp = await client.get(f"http://{COMFY_HOST}/history/{task_id}?token={token}", timeout=10.0)
+                if h_resp.status_code == 200:
+                    history_data = h_resp.json()
+                    if task_id in history_data:
+                        v_url = extract_video_url(history_data[task_id], token)
+                        complete_task(task_id, video_url=v_url)
+                        return # काम पूरा, लूप ख़तम!
+                        
+                # B. अगर नहीं बना, तो चेक करो कि अभी भी कतार में है क्या?
+                q_resp = await client.get(f"http://{COMFY_HOST}/queue?token={token}", timeout=10.0)
+                if q_resp.status_code == 200:
+                    queue_data = q_resp.json()
+                    
+                    is_running = any(req[1] == task_id for req in queue_data.get("queue_running", []))
+                    is_pending = any(req[1] == task_id for req in queue_data.get("queue_pending", []))
+                    
+                    if is_running:
+                        update_task_progress(task_id, progress_fake, "processing")
+                    elif is_pending:
+                        update_task_progress(task_id, progress_fake, "queued")
+                    else:
+                        # ना हिस्ट्री में है, ना रनिंग है, ना पेंडिंग है! मतलब सर्वर से क्रैश हो गया।
+                        complete_task(task_id, error="Task disappeared from ComfyUI server.")
+                        return
+                        
+        except Exception as e:
+            print(f"⚠️ Polling failed for {task_id}, retrying... ({e})")
+            
+        # 10 सेकंड आराम करके फिर चेक करेगा
+        await asyncio.sleep(10)
     try:
         async with httpx.AsyncClient() as client:
             h_resp = await client.get(f"http://{COMFY_HOST}/history/{task_id}?token={token}", timeout=30.0)