Update comfy_engine.py

comfy_engine.py

@@ -1,35 +1,39 @@
-import json
-import uuid
-import os
-import asyncio
+import json, uuid, os, asyncio, time
 import httpx
 import websockets
 from urllib.parse import urlencode
 
 COMFY_HOST = os.getenv("COMFY_HOST", "134.199.132.159")
 
-# वर्कफ़्लो लोड करें
 with open("workflow.json", "r", encoding="utf-8") as f:
     WORKFLOW_TEMPLATE = json.load(f)
 
+# 🧠 स्मार्ट रेसोल्यूशन डिक्शनरी (मल्टीपल्स ऑफ़ 16)
+RESOLUTIONS = {
+    "360p": {"16:9": (640, 360), "9:16": (360, 640), "1:1": (512, 512)},
+    "480p": {"16:9": (848, 480), "9:16": (480, 848), "1:1": (640, 640)},
+    "720p": {"16:9": (1280, 720), "9:16": (720, 1280), "1:1": (768, 768)}
+}
+
 def inject_params(req: dict) -> dict:
     p = json.loads(json.dumps(WORKFLOW_TEMPLATE))
-    # Text Prompts
+    
+    # प्रॉम्प्ट
     p["89"]["inputs"]["text"] = req.get("prompt", "")
-    if req.get("negative"):
-        p["72"]["inputs"]["text"] = req.get("negative")
     
-    # Video Settings
-    p["74"]["inputs"]["width"] = req.get("width", 640)
-    p["74"]["inputs"]["height"] = req.get("height", 640)
-    p["74"]["inputs"]["length"] = req.get("frames", 81)
-    p["88"]["inputs"]["fps"] = req.get("fps", 16)
+    # स्मार्ट डायमेंशन्स कैलकुलेशन
+    ratio = req.get("aspect_ratio", "16:9")
+    quality = req.get("quality", "480p")
+    width, height = RESOLUTIONS.get(quality, {}).get(ratio, (848, 480))
     
-    # Seed Configuration
-    seed = req.get("seed")
-    if seed:
-        p["81"]["inputs"]["noise_seed"] = seed
-        p["78"]["inputs"]["noise_seed"] = seed
+    p["74"]["inputs"]["width"] = width
+    p["74"]["inputs"]["height"] = height
+    
+    # फिक्स और हार्डकोडेड सेटिंग्स (बेस्ट क्वालिटी के लिए)
+    p["74"]["inputs"]["length"] = 81   # 81 फ्रेम्स
+    p["88"]["inputs"]["fps"] = 16      # 16 FPS (Smooth Slow-mo)
+    p["78"]["inputs"]["steps"] = 4     # 4 Steps
+    p["78"]["inputs"]["cfg"] = 5.0     # 5.0 CFG
     
     return p
 
@@ -41,8 +45,7 @@ def extract_video_url(history: dict, token: str) -> str:
                 it = node_out[key][0]
                 if all(k in it for k in ("filename", "subfolder", "type")):
                     q = urlencode(it)
-                    # 🔴 Proxy URL (No Mixed Content Error)
-                    return f"/api/video?{q}&token={token}"
+                    return f"/api/video?{q}&token={token}" # Proxy URL
     raise RuntimeError("No video file found in history")
 
 async def generate_video_stream(req: dict):
@@ -50,62 +53,55 @@ async def generate_video_stream(req: dict):
     client_id = str(uuid.uuid4())
     prompt = inject_params(req)
 
-    # 1. API को प्रॉम्प्ट भेजें
     async with httpx.AsyncClient() as client:
         resp = await client.post(f"http://{COMFY_HOST}/prompt?token={token}", json={"prompt": prompt, "client_id": client_id}, timeout=60.0)
         data = resp.json()
-        if "prompt_id" not in data:
-            yield f"data: {json.dumps({'error': 'Failed to queue prompt'})}\n\n"
+        prompt_id = data.get("prompt_id")
+        if not prompt_id:
+            yield f"data: {json.dumps({'error': 'Failed'})}\n\n"
             return
-        prompt_id = data["prompt_id"]
 
     total_nodes = max(1, len(prompt))
     seen = set()
+    start_time = time.time()
+    p_fake = 0
 
-    # 2. WebSocket से लाइव प्रोग्रेस ट्रैक करें (With Ping to keep alive)
     ws_url = f"ws://{COMFY_HOST}/ws?clientId={client_id}&token={token}"
     try:
-        async with websockets.connect(ws_url, ping_interval=20, ping_timeout=120) as ws:
+        async with websockets.connect(ws_url, ping_interval=20) as ws:
             while True:
                 out = await ws.recv()
+                
+                # 🚀 THE MAGIC HACK (Bytes Handling from your Gradio code)
+                if isinstance(out, (bytes, bytearray)):
+                    if p_fake < 95 and time.time() - start_time > 2:
+                        p_fake = min(95, p_fake + 1)
+                        yield f"data: {json.dumps({'progress': p_fake, 'status': 'rendering'})}\n\n"
+                    continue
+                
+                # Normal JSON handling
                 if isinstance(out, str):
                     msg = json.loads(out)
                     if msg.get("type") == "executing":
-                        data = msg.get("data", {})
-                        if data.get("prompt_id") != prompt_id:
-                            continue
-                        node = data.get("node")
+                        node = msg.get("data", {}).get("node")
                         if node is None:
-                            break # रेंडरिंग पूरी हो गई!
+                            break # डन!
                         if node not in seen:
                             seen.add(node)
-                            progress = min(99, int((len(seen) / total_nodes) * 100))
-                            yield f"data: {json.dumps({'progress': progress, 'status': 'rendering'})}\n\n"
+                            p_real = min(99, int((len(seen) / total_nodes) * 100))
+                            p_fake = max(p_fake, p_real) # जो भी ज़्यादा हो वो दिखाओ
+                            yield f"data: {json.dumps({'progress': p_fake, 'status': 'rendering'})}\n\n"
     except Exception as e:
-        # अगर तार टूट गया, तो एरर नहीं देंगे, बस "Finalizing" बोलेंगे
-        print(f"⚠️ WebSocket Disconnected: {e}. Switching to polling mode...")
-        yield f"data: {json.dumps({'progress': 99, 'status': 'Finalizing Video...'})}\n\n"
-
-    # 3. बैकअप प्लान (Fallback Polling) - अगर WS टूट गया तो इतिहास चेक करेंगे
-    max_retries = 60 # 60 * 3 = 180 seconds (3 mins max wait)
-    for _ in range(max_retries):
-        await asyncio.sleep(3) # 3 सेकंड रुकें
-        try:
-            async with httpx.AsyncClient() as client:
-                hist_resp = await client.get(f"http://{COMFY_HOST}/history/{prompt_id}?token={token}", timeout=30.0)
-                if hist_resp.status_code == 200:
-                    history_data = hist_resp.json()
-                    if prompt_id in history_data:
-                        history = history_data[prompt_id]
-                        try:
-                            final_url = extract_video_url(history, token)
-                            yield f"data: {json.dumps({'progress': 100, 'status': 'done', 'url': final_url})}\n\n"
-                            return # काम खतम!
-                        except Exception:
-                            pass # अभी वीडियो फाइल सेव हो रही है
-        except Exception as e:
-            print(f"Polling error: {e}")
-            pass # नेटवर्क ग्लिच, इग्नोर करें और फिर ट्राई करें
+        print(f"WS Disconnected: {e}")
 
-    yield f"data: {json.dumps({'error': 'Timeout: Rendering took too long.'})}\n\n"
-    
+    # हिस्ट्री से वीडियो निकालें
+    await asyncio.sleep(2)
+    async with httpx.AsyncClient() as client:
+        hist_resp = await client.get(f"http://{COMFY_HOST}/history/{prompt_id}?token={token}", timeout=60.0)
+        history = hist_resp.json().get(prompt_id, {})
+    
+    try:
+        final_url = extract_video_url(history, token)
+        yield f"data: {json.dumps({'progress': 100, 'status': 'done', 'url': final_url})}\n\n"
+    except Exception as e:
+        yield f"data: {json.dumps({'error': str(e)})}\n\n"