Update app.py

app.py

@@ -1,9 +1,14 @@
 import asyncio
 import json
 import httpx
+import logging
 from fastapi import FastAPI, BackgroundTasks
 from pydantic import BaseModel
 
+# تنظیمات لاگینگ برای دیدن خطاها در کنسول
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
 app = FastAPI()
 
 # آدرس اصلی API مدل Wan2.2
@@ -16,39 +21,73 @@ class TaskPayload(BaseModel):
     resolution: str
     callback_url: str
 
+def extract_video_url(data):
+    """تابع هوشمند برای بیرون کشیدن لینک ویدیو از هر ساختار جیسونی"""
+    try:
+        # اگر دیتا لیست است، اولین آیتم را چک کن
+        if isinstance(data, list) and len(data) > 0:
+            item = data[0]
+        else:
+            item = data
+
+        # 1. چک کردن ساختار استاندارد Gradio Video
+        if isinstance(item, dict):
+            if "video" in item and isinstance(item["video"], dict):
+                return item["video"].get("url") or item["video"].get("path")
+            if "url" in item:
+                return item["url"]
+            if "path" in item:
+                return item["path"]
+            if "name" in item:
+                return f"/file={item['name']}"
+        
+        # 2. چک کردن رشته ساده
+        if isinstance(item, str):
+            if item.endswith(".mp4") or "/file=" in item:
+                return item
+                
+    except Exception as e:
+        logger.error(f"Error extracting URL: {e}")
+    return None
+
 async def upload_file_from_url(client: httpx.AsyncClient, file_url: str, file_type: str):
     """دانلود فایل از مدیر و آپلود به Wan2.2"""
-    print(f"Downloading {file_type} from {file_url}...")
-    # 1. دانلود فایل از اسپیس مدیر
-    file_resp = await client.get(file_url)
+    logger.info(f"Downloading {file_type} from {file_url}...")
+    
+    # دانلود با ریدایرکت و تایم‌اوت مناسب
+    file_resp = await client.get(file_url, follow_redirects=True)
     if file_resp.status_code != 200:
         raise Exception(f"Failed to download {file_type}: {file_resp.status_code}")
     
     file_bytes = file_resp.content
     filename = file_url.split("/")[-1]
     
-    # 2. آپلود به Wan2.2
-    print(f"Uploading {filename} to Wan2.2...")
+    # آپلود به Wan2.2
+    logger.info(f"Uploading {filename} to Wan2.2...")
     files = {'files': (filename, file_bytes)}
     upload_resp = await client.post(f"{WAN_API_BASE}/upload", files=files)
     
     if upload_resp.status_code != 200:
-        raise Exception(f"Failed to upload {file_type} to Wan2.2")
+        raise Exception(f"Failed to upload {file_type} to Wan2.2: {upload_resp.text}")
     
+    # سرور لیست برمی‌گرداند: ["/tmp/..."]
     remote_path = upload_resp.json()[0]
-    print(f"Uploaded: {remote_path}")
+    logger.info(f"Uploaded {file_type}: {remote_path}")
     return remote_path
 
 async def process_task(payload: TaskPayload):
-    print(f"Processing Job {payload.job_id}")
+    logger.info(f"--- Processing Job {payload.job_id} ---")
     
-    async with httpx.AsyncClient(timeout=120.0) as client: # تایم‌اوت بالا
+    # تایم‌اوت کلی برای کل عملیات (۳۰۰ ثانیه = ۵ دقیقه)
+    timeout = httpx.Timeout(300.0, connect=60.0)
+    
+    async with httpx.AsyncClient(timeout=timeout) as client:
         try:
             # الف) آپلود فایل‌ها
             img_remote = await upload_file_from_url(client, payload.image_url, "image")
             aud_remote = await upload_file_from_url(client, payload.audio_url, "audio")
             
-            # ب) ارسال درخواست Predict
+            # ب) ارسال درخواست ساخت
             req_data = {
                 "data": [
                     {"path": img_remote, "meta": {"_type": "gradio.FileData"}},
@@ -57,65 +96,89 @@ async def process_task(payload: TaskPayload):
                 ]
             }
             
+            logger.info("Sending predict request...")
             predict_resp = await client.post(f"{WAN_API_BASE}/call/predict", json=req_data)
+            
             if predict_resp.status_code != 200:
                 raise Exception(f"Predict call failed: {predict_resp.text}")
             
             event_id = predict_resp.json().get("event_id")
-            print(f"Event ID: {event_id}")
+            logger.info(f"Task submitted. Event ID: {event_id}")
             
-            # ج) گوش دادن به استریم (Polling the stream)
-            # در پایتون، به جای EventSource، می‌توانیم با یک لوپ ساده GET بزنیم یا stream کنیم
+            # ج) گوش دادن به استریم
             final_video_url = None
+            stream_url = f"{WAN_API_BASE}/call/predict/{event_id}"
             
-            async with client.stream("GET", f"{WAN_API_BASE}/call/predict/{event_id}", headers={"Accept": "text/event-stream"}) as response:
+            # تایم‌اوت برای خواندن استریم باید خیلی زیاد باشد
+            async with client.stream("GET", stream_url, headers={"Accept": "text/event-stream"}, timeout=300.0) as response:
                 async for line in response.aiter_lines():
+                    # خط‌های خالی را رد کن
+                    if not line.strip():
+                        continue
+                    
+                    # لاگ کردن خطوط مهم برای دیباگ (اختیاری)
+                    # logger.info(f"Stream: {line[:100]}") 
+
+                    if line.startswith("event: error"):
+                         raise Exception("Server sent an error event.")
+
                     if line.startswith("data: "):
                         try:
-                            data = json.loads(line[6:])
-                            # بررسی ساختار خروجی
-                            if isinstance(data, list) and len(data) > 0:
-                                result_item = data[0]
-                                # پیدا کردن URL در ساختارهای مختلف
-                                if isinstance(result_item, dict):
-                                    if "video" in result_item and "url" in result_item["video"]:
-                                        final_video_url = result_item["video"]["url"]
-                                    elif "url" in result_item:
-                                        final_video_url = result_item["url"]
-                                    elif "name" in result_item:
-                                        final_video_url = f"https://wan-ai-wan2-2-s2v.ms.show/file={result_item['name']}"
-                                elif isinstance(result_item, str) and (result_item.endswith(".mp4") or "/file=" in result_item):
-                                     final_video_url = result_item
+                            data_str = line[6:].strip()
+                            # داده‌های خالی یا null را رد کن
+                            if not data_str or data_str == "null":
+                                continue
                                 
-                                if final_video_url:
-                                    # اصلاح لینک نسبی
-                                    if final_video_url.startswith("/"):
-                                        final_video_url = f"https://wan-ai-wan2-2-s2v.ms.show{final_video_url}"
-                                    break
-                        except:
+                            data = json.loads(data_str)
+                            
+                            # تلاش برای یافتن ویدیو
+                            found = extract_video_url(data)
+                            if found:
+                                final_video_url = found
+                                logger.info(f"Video URL found: {final_video_url}")
+                                # اینجا break نمیکنیم، شاید پیام complete بعدش بیاد که مطمئن بشیم
+                        except json.JSONDecodeError:
                             pass
+                        except Exception as e:
+                            logger.error(f"Parsing error: {e}")
 
             if final_video_url:
-                print(f"Success! Video: {final_video_url}")
-                # د) بازگرداندن نتیجه به مدیر
+                # اصلاح لینک اگر نسبی باشد
+                if final_video_url.startswith("/"):
+                    final_video_url = f"https://wan-ai-wan2-2-s2v.ms.show{final_video_url}"
+                
+                # اصلاح لینک‌های عجیب //file=
+                if "//file=" in final_video_url:
+                    final_video_url = final_video_url.replace("//file=", "/file=")
+
+                logger.info(f"SUCCESS! Sending callback: {final_video_url}")
+                
+                # د) ارسال به مدیر
                 await client.post(payload.callback_url, json={
                     "job_id": payload.job_id,
                     "status": "COMPLETED",
                     "video_url": final_video_url
                 })
             else:
-                raise Exception("Stream ended without video URL")
+                raise Exception("Stream ended but no video URL was found in data.")
 
         except Exception as e:
-            print(f"Task Failed: {e}")
-            await client.post(payload.callback_url, json={
-                "job_id": payload.job_id,
-                "status": "FAILED",
-                "message": str(e)
-            })
+            logger.error(f"Task Failed: {e}")
+            try:
+                await client.post(payload.callback_url, json={
+                    "job_id": payload.job_id,
+                    "status": "FAILED",
+                    "message": str(e)
+                })
+            except Exception as cb_e:
+                logger.error(f"Callback failed: {cb_e}")
 
 @app.post("/process")
 async def accept_task(payload: TaskPayload, background_tasks: BackgroundTasks):
-    """تسک را می‌پذیرد و بلافاصله OK برمی‌گرداند تا تایم‌اوت نشود"""
+    logger.info(f"Received task {payload.job_id}")
     background_tasks.add_task(process_task, payload)
-    return {"status": "accepted"}
+    return {"status": "accepted"}
+
+@app.get("/")
+async def root():
+    return {"status": "Worker is running", "target": WAN_API_BASE}