Update app.py

app.py

@@ -1,18 +1,25 @@
-import asyncio
-import json
-import httpx
+import os
+import shutil
 import logging
+import uuid
+import asyncio
 from fastapi import FastAPI, BackgroundTasks
 from pydantic import BaseModel
+from gradio_client import Client, handle_file
+import httpx
 
-# تنظیمات لاگینگ برای دیدن خطاها در کنسول
+# تنظیمات لاگ
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
-# آدرس اصلی API مدل Wan2.2
-WAN_API_BASE = "https://wan-ai-wan2-2-s2v.ms.show/gradio_api"
+# آدرس سرور هوش مصنوعی (بدون /gradio_api)
+HF_MODEL_URL = "https://wan-ai-wan2-2-s2v.ms.show/"
+
+# پوشه موقت برای دانلود فایل‌ها
+TEMP_DIR = "/tmp/worker_files"
+os.makedirs(TEMP_DIR, exist_ok=True)
 
 class TaskPayload(BaseModel):
     job_id: str
@@ -21,164 +28,212 @@ 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"""
-    logger.info(f"Downloading {file_type} from {file_url}...")
+async def download_file(url: str, suffix: str) -> str:
+    """فایل را از مدیر دانلود و در پوشه موقت ذخیره می‌کند"""
+    local_filename = f"{uuid.uuid4()}.{suffix}"
+    local_path = os.path.join(TEMP_DIR, local_filename)
     
-    # دانلود با ریدایرکت و تایم‌اوت مناسب
-    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}")
+    logger.info(f"Downloading {url} to {local_path}...")
     
-    file_bytes = file_resp.content
-    filename = file_url.split("/")[-1]
+    async with httpx.AsyncClient() as client:
+        resp = await client.get(url, follow_redirects=True, timeout=60.0)
+        if resp.status_code != 200:
+            raise Exception(f"Download failed: {resp.status_code}")
+        
+        with open(local_path, "wb") as f:
+            f.write(resp.content)
+            
+    return local_path
+
+def run_gradio_client(img_path, aud_path, resolution):
+    """این تابع به صورت همگام (Sync) اجرا می‌شود چون کتابخانه Gradio همگام است"""
+    logger.info("Connecting to Gradio Client...")
+    client = Client(HF_MODEL_URL)
     
-    # آپلود به 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)
+    logger.info("Sending request to Wan2.2 (This may take minutes)...")
     
-    if upload_resp.status_code != 200:
-        raise Exception(f"Failed to upload {file_type} to Wan2.2: {upload_resp.text}")
+    # ارسال درخواست و انتظار برای نتیجه (Blocking)
+    result = client.predict(
+        ref_img=handle_file(img_path),
+        audio=handle_file(aud_path),
+        resolution=resolution,
+        api_name="/predict"
+    )
     
-    # سرور لیست برمی‌گرداند: ["/tmp/..."]
-    remote_path = upload_resp.json()[0]
-    logger.info(f"Uploaded {file_type}: {remote_path}")
-    return remote_path
+    # نتیجه معمولاً یک لیست است: [مسیر_ویدیو, پیام_وضعیت]
+    # یا گاهی یک دیکشنری بسته به نسخه
+    logger.info(f"Raw Result received: {result}")
+    return result
 
 async def process_task(payload: TaskPayload):
     logger.info(f"--- Processing Job {payload.job_id} ---")
+    img_path = None
+    aud_path = None
     
-    # تایم‌اوت کلی برای کل عملیات (۳۰۰ ثانیه = ۵ دقیقه)
-    timeout = httpx.Timeout(300.0, connect=60.0)
+    try:
+        # 1. دانلود فایل‌ها از سرور مدیر به ��ضای لوکال کارگر
+        img_path = await download_file(payload.image_url, "png") # پسوند حدسی
+        aud_path = await download_file(payload.audio_url, "mp3")
+        
+        # 2. اجرای کلاینت Gradio در یک Thread جداگانه (تا Event Loop اصلی بلاک نشود)
+        # این خط معجزه می‌کند: صبر می‌کند تا ویدیو واقعاً ساخته شود
+        result = await asyncio.to_thread(
+            run_gradio_client, 
+            img_path, 
+            aud_path, 
+            payload.resolution
+        )
+        
+        # 3. استخراج لینک ویدیو
+        final_video_url = None
+        
+        # بررسی فرمت خروجی Gradio Client
+        if isinstance(result, list) and len(result) > 0:
+            # در نسخه کلاینت، فایل خروجی دانلود شده و در تمپ لوکال است
+            # اما ما نیاز به لینک عمومی داریم.
+            # نکته: کلاینت فایل را دانلود میکند. ما باید لینک اصلی را پیدا کنیم 
+            # یا فایل دانلود شده را جایی آپلود کنیم.
+            # راه ساده‌تر برای این سناریو: استفاده از لینک Gradio API
+            
+            video_data = result[0]
+            if isinstance(video_data, dict) and "url" in video_data:
+                final_video_url = video_data["url"]
+            elif isinstance(video_data, str) and os.path.exists(video_data):
+                # اگر فایل دانلود شده است، یعنی موفقیت آمیز بوده.
+                # اما ما نمیتوانیم فایل لوکال کارگر را به مدیر بدهیم (چون سرور جداست)
+                # مگر اینکه کارگر آن را آپلود کند.
+                # خوشبختانه Gradio Client معمولا آبجکت اصلی را هم برمیگرداند.
+                pass
+
+        # از آنجا که ما از طریق API کار میکنیم، فایل روی سرور ModelScope ساخته شده.
+        # کلاینت پایتون آن را دانلود میکند.
+        # برای اینکه لینک قابل دانلود به کاربر بدهیم، باید کمی کلک بزنیم:
+        # چون فایل به Worker دانلود شده، ما نمیتوانیم لینک لوکال بدهیم.
+        # راه حل: فایل تولید شده در Worker را به Manager آپلود کنیم؟ نه پیچیده است.
+        # راه حل بهتر: استفاده از خروجی خام بدون دانلود خودکار (که سخت است).
+        
+        # بیایید فرض کنیم کلاینت فایل را دانلود کرده است `video_path`
+        local_video_result = result[0] # مسیر فایل ویدیوی تولید شده روی دیسک کارگر
+        
+        if local_video_result and os.path.exists(local_video_result):
+            logger.info(f"Video generated locally at: {local_video_result}")
+            
+            # حالا باید این ویدیو را به سرور مدیر برگردانیم
+            # چون لینک مستقیم نداریم (فایل پرایوت است)، آن را به مدیر آپلود میکنیم
+            
+            async with httpx.AsyncClient(timeout=60.0) as client:
+                with open(local_video_result, "rb") as f:
+                    files = {"file": ("generated_video.mp4", f, "video/mp4")}
+                    # فرض میکنیم مدیر یک اندپوینت برای دریافت فایل نهایی دارد
+                    # اگر ندارد، فعلا لینک تسک را میفرستیم (که با این روش کار نمیکند)
+                    pass
+            
+            # --- اصلاح استراتژی برای سادگی ---
+            # چون انتقال فایل بین سرورها سخت است، ما لینک Gradio را میخواهیم.
+            # اما gradio_client فایل را دانلود میکند.
+            # بیایید دوباره به روش httpx برگردیم اما با استراتژی Poll (نه Stream)
+            raise Exception("Switching logic") # Jump to except block to retry logic? No.
+
+    except Exception as e:
+        pass
+
+    # ==================================================================
+    # روش دوم و نهایی (ترکیبی):
+    # استفاده از gradio_client فقط برای آپلود و predict، اما گرفتن لینک وب
+    # ==================================================================
     
-    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")
+    try:
+        # تلاش مجدد با کلاینت اما استخراج اطلاعات متا
+        # متاسفانه gradio_client فایل را دانلود میکند و لینک اصلی را مخفی میکند.
+        # پس برمیگردیم به httpx اما بدون stream. 
+        # به جای stream از polling استفاده میکنیم که قطع نمیشود.
+        
+        async with httpx.AsyncClient(timeout=300.0) as client:
+             # 1. آپلود (همان کد قبلی که کار میکرد)
+            logger.info("Uploading files manually to ensure we get URL...")
+            
+            # دانلود فایل از مدیر
+            f_img = await client.get(payload.image_url)
+            f_aud = await client.get(payload.audio_url)
+            
+            # آپلود به Gradio
+            up_img = await client.post(f"{WAN_API_BASE}/upload", files={'files': ('img.png', f_img.content)})
+            up_aud = await client.post(f"{WAN_API_BASE}/upload", files={'files': ('aud.mp3', f_aud.content)})
+            
+            remote_img = up_img.json()[0]
+            remote_aud = up_aud.json()[0]
             
-            # ب) ارسال درخواست ساخت
+            # 2. Submit Request
             req_data = {
                 "data": [
-                    {"path": img_remote, "meta": {"_type": "gradio.FileData"}},
-                    {"path": aud_remote, "meta": {"_type": "gradio.FileData"}},
+                    {"path": remote_img, "meta": {"_type": "gradio.FileData"}},
+                    {"path": remote_aud, "meta": {"_type": "gradio.FileData"}},
                     payload.resolution
                 ]
             }
             
-            logger.info("Sending predict request...")
-            predict_resp = await client.post(f"{WAN_API_BASE}/call/predict", json=req_data)
+            resp = await client.post(f"{WAN_API_BASE}/call/predict", json=req_data)
+            event_id = resp.json()['event_id']
+            logger.info(f"Event ID: {event_id} - Polling started")
             
-            if predict_resp.status_code != 200:
-                raise Exception(f"Predict call failed: {predict_resp.text}")
-            
-            event_id = predict_resp.json().get("event_id")
-            logger.info(f"Task submitted. Event ID: {event_id}")
-            
-            # ج) گوش دادن به استریم
-            final_video_url = None
-            stream_url = f"{WAN_API_BASE}/call/predict/{event_id}"
-            
-            # تایم‌اوت برای خواندن استریم باید خیلی زیاد باشد
-            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_str = line[6:].strip()
-                            # داده‌های خالی یا null را رد کن
-                            if not data_str or data_str == "null":
-                                continue
-                                
-                            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:
-                # اصلاح لینک اگر نسبی باشد
-                if final_video_url.startswith("/"):
-                    final_video_url = f"https://wan-ai-wan2-2-s2v.ms.show{final_video_url}"
+            # 3. POLLING (به جای Streaming)
+            # هر 3 ثانیه وضعیت را چک میکنیم
+            final_url = None
+            for i in range(100): # تا 300 ثانیه (5 دقیقه) تلاش کن
+                await asyncio.sleep(3)
                 
-                # اصلاح لینک‌های عجیب //file=
-                if "//file=" in final_video_url:
-                    final_video_url = final_video_url.replace("//file=", "/file=")
+                # در Gradio 5 برای گ��فتن وضعیت باید به استریم وصل شد، اما اگر قطع شد مهم نیست
+                # راه حل جایگزین: استفاده از requests معمولی (نه stream) به صورت تک‌ضرب
+                # متاسفانه API استاندارد polling ندارد. 
+                
+                # بیایید دوباره Stream را تست کنیم اما با تنظیمات بسیار خاص که قطع نشود
+                try:
+                    async with client.stream("GET", f"{WAN_API_BASE}/call/predict/{event_id}", headers={"Accept": "text/event-stream"}, timeout=10.0) as stream_resp:
+                        async for line in stream_resp.aiter_lines():
+                            if line.startswith("data: "):
+                                try:
+                                    data = json.loads(line[6:])
+                                    # لاگ کردن دیتا برای فهمیدن وضعیت
+                                    logger.info(f"Poll Status: {str(data)[:50]}...") 
+                                    
+                                    if isinstance(data, list) and len(data) > 0:
+                                        # چک کردن موفقیت
+                                        if data[0] and "video" in str(data[0]):
+                                            # استخراج لینک
+                                            res = data[0]
+                                            if isinstance(res, dict) and "video" in res:
+                                                final_url = res["video"]["url"]
+                                            elif isinstance(res, dict) and "url" in res:
+                                                final_url = res["url"]
+                                            
+                                            if final_url:
+                                                break # Loop line
+                                except:
+                                    pass
+                        if final_url: break # Loop poll
+                except Exception as stream_err:
+                    logger.warning(f"Stream checking interrupted, retrying... {stream_err}")
+                    continue
 
-                logger.info(f"SUCCESS! Sending callback: {final_video_url}")
+            if final_url:
+                if final_url.startswith("/"):
+                    final_url = f"https://wan-ai-wan2-2-s2v.ms.show{final_url}"
                 
-                # د) ارسال به مدیر
+                logger.info(f"✅ SUCCESS! Video: {final_url}")
                 await client.post(payload.callback_url, json={
-                    "job_id": payload.job_id,
-                    "status": "COMPLETED",
-                    "video_url": final_video_url
+                    "job_id": payload.job_id, "status": "COMPLETED", "video_url": final_url
                 })
             else:
-                raise Exception("Stream ended but no video URL was found in data.")
+                raise Exception("Timed out waiting for video.")
 
-        except Exception as 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}")
+    except Exception as e:
+        logger.error(f"FAILED: {e}")
+        async with httpx.AsyncClient() as client:
+            await client.post(payload.callback_url, json={
+                "job_id": payload.job_id, "status": "FAILED", "message": str(e)
+            })
 
 @app.post("/process")
 async def accept_task(payload: TaskPayload, background_tasks: BackgroundTasks):
-    logger.info(f"Received task {payload.job_id}")
     background_tasks.add_task(process_task, payload)
-    return {"status": "accepted"}
-
-@app.get("/")
-async def root():
-    return {"status": "Worker is running", "target": WAN_API_BASE}
+    return {"status": "accepted"}