Update app.py

app.py

@@ -1,12 +1,10 @@
-import os
-import shutil
-import logging
-import uuid
 import asyncio
+import json
+import httpx
+import logging
+import time
 from fastapi import FastAPI, BackgroundTasks
 from pydantic import BaseModel
-from gradio_client import Client, handle_file
-import httpx
 
 # تنظیمات لاگ
 logging.basicConfig(level=logging.INFO)
@@ -14,12 +12,9 @@ logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
-# آدرس سرور هوش مصنوعی (بدون /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)
+# --- تنظیمات حیاتی ---
+# این متغیر که قبلاً خطا می‌داد الان تعریف شده است
+WAN_API_BASE = "https://wan-ai-wan2-2-s2v.ms.show/gradio_api"
 
 class TaskPayload(BaseModel):
     job_id: str
@@ -28,212 +23,144 @@ class TaskPayload(BaseModel):
     resolution: str
     callback_url: str
 
-async def download_file(url: str, suffix: str) -> str:
-    """فایل را از مدیر دانلود و در پوشه موقت ذخیره می‌کند"""
-    local_filename = f"{uuid.uuid4()}.{suffix}"
-    local_path = os.path.join(TEMP_DIR, local_filename)
-    
-    logger.info(f"Downloading {url} to {local_path}...")
-    
-    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)
-    
-    logger.info("Sending request to Wan2.2 (This may take minutes)...")
+async def upload_to_wan(client: httpx.AsyncClient, file_url: str):
+    """فایل را از مدیر دانلود و مستقیم به سرور هوش مصنوعی آپلود می‌کند"""
+    filename = file_url.split("/")[-1]
+    logger.info(f"Downloading {filename}...")
     
-    # ارسال درخواست و انتظار برای نتیجه (Blocking)
-    result = client.predict(
-        ref_img=handle_file(img_path),
-        audio=handle_file(aud_path),
-        resolution=resolution,
-        api_name="/predict"
-    )
+    # دانلود از سرور خودمان
+    resp = await client.get(file_url, timeout=60.0)
+    if resp.status_code != 200:
+        raise Exception(f"Download failed: {resp.status_code}")
     
-    # نتیجه معمولاً یک لیست است: [مسیر_ویدیو, پیام_وضعیت]
-    # یا گاهی یک دیکشنری بسته به نسخه
-    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
+    file_bytes = resp.content
     
+    # آپلود به سرور Wan
+    logger.info(f"Uploading {filename} to External AI Server...")
+    files = {'files': (filename, file_bytes)}
     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
-        )
+        # تلاش برای آپلود (ممکن است چند بار نیاز باشد)
+        wan_resp = await client.post(f"{WAN_API_BASE}/upload", files=files, timeout=60.0)
+        if wan_resp.status_code != 200:
+            raise Exception(f"Upload failed: {wan_resp.text}")
         
-        # 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.
-
+        remote_path = wan_resp.json()[0]
+        return remote_path
     except Exception as e:
-        pass
+        raise Exception(f"Upload connection error: {e}")
 
-    # ==================================================================
-    # روش دوم و نهایی (ترکیبی):
-    # استفاده از gradio_client فقط برای آپلود و predict، اما گرفتن لینک وب
-    # ==================================================================
+async def process_task(payload: TaskPayload):
+    logger.info(f"=== Starting Job {payload.job_id} ===")
     
-    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)
+    async with httpx.AsyncClient(timeout=None) as client: # تایم‌اوت کلی غیرفعال
+        try:
+            # 1. آماده‌سازی فایل‌ها
+            img_remote = await upload_to_wan(client, payload.image_url)
+            aud_remote = await upload_to_wan(client, 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
+            # 2. ثبت درخواست (Predict)
             req_data = {
                 "data": [
-                    {"path": remote_img, "meta": {"_type": "gradio.FileData"}},
-                    {"path": remote_aud, "meta": {"_type": "gradio.FileData"}},
+                    {"path": img_remote, "meta": {"_type": "gradio.FileData"}},
+                    {"path": aud_remote, "meta": {"_type": "gradio.FileData"}},
                     payload.resolution
                 ]
             }
             
-            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")
+            logger.info("Submitting prediction task...")
+            predict_resp = await client.post(f"{WAN_API_BASE}/call/predict", json=req_data, timeout=30.0)
+            if predict_resp.status_code != 200:
+                raise Exception(f"Prediction failed: {predict_resp.text}")
             
-            # 3. POLLING (به جای Streaming)
-            # هر 3 ثانیه وضعیت را چک میکنیم
-            final_url = None
-            for i in range(100): # تا 300 ثانیه (5 دقیقه) تلاش کن
-                await asyncio.sleep(3)
-                
-                # در Gradio 5 برای گرفتن وضعیت باید به استریم وصل شد، اما اگر قطع شد مهم نیست
-                # راه حل جایگزین: استفاده از requests معمولی (نه stream) به صورت تک‌ضرب
-                # متاسفانه API استاندارد polling ندارد. 
-                
-                # بیایید دوباره Stream را تست کنیم اما با تنظیمات بسیار خاص که قطع نشود
+            event_id = predict_resp.json().get("event_id")
+            logger.info(f"Job Queued! Event ID: {event_id}")
+            
+            # 3. حلقه انتظار هوشمند (Robust Polling)
+            # این حلقه تا 10 دقیقه (600 ثانیه) تلاش می‌کند
+            start_time = time.time()
+            final_video_url = None
+            
+            while time.time() - start_time < 600:
                 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():
+                    # اتصال به استریم
+                    stream_url = f"{WAN_API_BASE}/call/predict/{event_id}"
+                    async with client.stream("GET", stream_url, headers={"Accept": "text/event-stream"}, timeout=20.0) as response:
+                        async for line in response.aiter_lines():
+                            if not line.strip(): continue
+                            
                             if line.startswith("data: "):
                                 try:
                                     data = json.loads(line[6:])
-                                    # لاگ کردن دیتا برای فهمیدن وضعیت
-                                    logger.info(f"Poll Status: {str(data)[:50]}...") 
                                     
+                                    # لاگ وضعیت (اختیاری)
+                                    # logger.info(f"Status check: {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"]
+                                        result = data[0]
+                                        # الگوهای مختلف پیدا کردن لینک
+                                        found_url = None
+                                        if isinstance(result, dict):
+                                            found_url = result.get("video", {}).get("url") or result.get("url") or (f"/file={result['name']}" if "name" in result else None)
+                                        elif isinstance(result, str) and ("/file=" in result or ".mp4" in result):
+                                            found_url = result
+                                        
+                                        if found_url:
+                                            final_video_url = found_url
+                                            break # شکستن حلقه خواندن خط
                                             
-                                            if final_url:
-                                                break # Loop line
-                                except:
+                                except Exception:
                                     pass
-                        if final_url: break # Loop poll
-                except Exception as stream_err:
-                    logger.warning(f"Stream checking interrupted, retrying... {stream_err}")
-                    continue
+                        
+                    if final_video_url:
+                        break # شکستن حلقه اصلی زمان
+                        
+                    # اگر استریم قطع شد ولی ویدیو هنوز آماده نیست، کمی صبر کن و دوباره وصل شو
+                    logger.info("Stream disconnected or empty, reconnecting to check status...")
+                    await asyncio.sleep(3)
+
+                except Exception as e:
+                    logger.warning(f"Connection glitch ({e}), retrying in 3 seconds...")
+                    await asyncio.sleep(3)
+
+            # 4. پایان کار
+            if final_video_url:
+                # اصلاح لینک
+                if final_video_url.startswith("/"):
+                    final_video_url = f"https://wan-ai-wan2-2-s2v.ms.show{final_video_url}"
+                
+                # فیکس کردن باگ دو اسلش
+                final_video_url = final_video_url.replace("//file=", "/file=")
 
-            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 URL: {final_video_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_url
+                    "job_id": payload.job_id,
+                    "status": "COMPLETED",
+                    "video_url": final_video_url
                 })
             else:
-                raise Exception("Timed out waiting for video.")
+                raise Exception("Timeout: Video processing took too long (>10 mins)")
 
-    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)
-            })
+        except Exception as e:
+            logger.error(f"❌ Job Failed: {e}")
+            # تلاش برای اعلام خطا به مدیر
+            try:
+                await client.post(payload.callback_url, json={
+                    "job_id": payload.job_id,
+                    "status": "FAILED",
+                    "message": str(e)
+                })
+            except:
+                pass
 
 @app.post("/process")
 async def accept_task(payload: TaskPayload, background_tasks: BackgroundTasks):
+    # تسک را می‌پذیریم و در پس‌زمینه اجرا می‌کنیم
     background_tasks.add_task(process_task, payload)
-    return {"status": "accepted"}
+    return {"status": "accepted"}
+
+@app.get("/")
+async def root():
+    return {"status": "Worker Ready", "api_target": WAN_API_BASE}