Update app.py

app.py

@@ -1,116 +1,261 @@
 """
-JoyAI-Video-Edit — ZeroGPU Compatible Entry Point
-این فایل به عنوان نقطه ورود برای HF Spaces طراحی شده است.
-از آنجا که JoyAI اصلی از ماژول داخلی xvideo استفاده می‌کند،
-این اسکریپت یک رابط Gradio/ZeroGPU سازگار فراهم می‌کند.
+JoyAI-Video-Edit — Complete ZeroGPU + Gradio Entry Point
+شامل: دانلود خودکار رانتایم + چک‌پوینت‌ها + رابط کاربری Gradio
 """
 import os
 import sys
+import subprocess
 from pathlib import Path
 
-# ✅ تنظیم مسیرهای قابل نوشتن در HF Spaces (رفع PermissionError)
+# ============================================================
+# ✅ ۱. تنظیم مسیرهای قابل نوشتن (رفع PermissionError)
+# ============================================================
 os.environ["HF_HOME"] = "/tmp/huggingface"
 os.environ["XDG_CACHE_HOME"] = "/tmp/cache"
 os.environ["JOYOMNI_CKPT_ROOT"] = "/tmp/checkpoints/JoyAI-Video-Edit"
-os.environ["COMFYUI_MODEL_PATH"] = "/tmp/comfyui_models"
+os.environ["JOYAI_REPO"] = "/tmp/JoyAI-Video-Edit"
 
-# ایجاد دایرکتوری‌ها قبل از هرگونه import
-for p in [
+REQUIRED_DIRS = [
     os.environ["HF_HOME"],
     os.environ["JOYOMNI_CKPT_ROOT"],
-    os.environ["COMFYUI_MODEL_PATH"],
-    "/tmp/comfyui_models/diffusion_models",
-]:
-    Path(p).mkdir(parents=True, exist_ok=True)
+    os.environ["JOYAI_REPO"],
+    f"{os.environ['JOYOMNI_CKPT_ROOT']}/dit",
+    f"{os.environ['JOYOMNI_CKPT_ROOT']}/vae",
+    f"{os.environ['JOYOMNI_CKPT_ROOT']}/text_encoder",
+]
+for d in REQUIRED_DIRS:
+    Path(d).mkdir(parents=True, exist_ok=True)
 
-# ✅ الزامی برای شناسایی توسط ZeroGPU
+print("📁 Writable directories initialized in /tmp/")
+
+
+# ============================================================
+# ✅ ۲. دانلود خودکار رانتایم xvideo از گیت‌هاب
+# ============================================================
+def setup_runtime():
+    """کلون کردن و نصب رانتایم JoyAI-Video-Edit"""
+    repo_path = Path(os.environ["JOYAI_REPO"])
+    
+    if not (repo_path / "deploy").exists():
+        print("📥 Cloning JoyAI-Video-Edit runtime from GitHub...")
+        result = subprocess.run(
+            ["git", "clone", "--depth", "1",
+             "https://github.com/jd-opensource/JoyAI-Video-Edit.git",
+             str(repo_path)],
+            capture_output=True, text=True
+        )
+        if result.returncode != 0:
+            print(f"❌ Git clone failed: {result.stderr}")
+            return False
+        print("✅ Runtime cloned successfully.")
+    else:
+        print("ℹ️ Runtime already exists, skipping clone.")
+    
+    # اضافه کردن deploy به Python path
+    deploy_path = str(repo_path / "deploy")
+    if deploy_path not in sys.path:
+        sys.path.insert(0, deploy_path)
+        print(f"🔗 Added to sys.path: {deploy_path}")
+    
+    return True
+
+
+# ============================================================
+# ✅ ۳. دانلود خودکار چک‌پوینت‌ها از HuggingFace
+# ============================================================
+def download_checkpoints():
+    """دانلود وزن‌های مدل از HF Hub"""
+    ckpt_root = Path(os.environ["JOYOMNI_CKPT_ROOT"])
+    
+    # بررسی اینکه آیا چک‌پوینت‌ها قبلاً دانلود شده‌اند
+    dit_path = ckpt_root / "dit" / "joyai_video_edit_dit_0811.pth"
+    vae_config = ckpt_root / "vae" / "config.json"
+    
+    if dit_path.exists() and vae_config.exists():
+        print("ℹ️ Checkpoints already downloaded.")
+        return True
+    
+    try:
+        from huggingface_hub import hf_hub_download, snapshot_download
+        
+        print("📥 Downloading JoyAI-Video-Edit checkpoints (~30GB)...")
+        print("   This may take 10-20 minutes on first run.")
+        
+        # دانلود کل پوشه چک‌پوینت
+        snapshot_download(
+            repo_id="jdopensource/JoyAI-Video-Edit",
+            local_dir=str(ckpt_root),
+            allow_patterns=["dit/**", "vae/**", "text_encoder/**"],
+            ignore_patterns=["*.md", "*.txt", ".git*"],
+        )
+        
+        print("✅ Checkpoints downloaded successfully.")
+        return True
+        
+    except Exception as e:
+        print(f"❌ Checkpoint download failed: {e}")
+        print("💡 You may need to manually upload checkpoints to this Space.")
+        return False
+
+
+# ============================================================
+# ✅ ۴. بارگذاری مدل روی ZeroGPU
+# ============================================================
 import spaces
 
-@spaces.GPU(duration=60, gpu_size="xlarge")
+runtime = None
+load_error = None
+
+@spaces.GPU(duration=120, gpu_size="xlarge")
 def load_model():
-    """بارگذاری مدل در GPU با مدیریت خطا"""
+    """بارگذاری JoyAI-Video-Edit روی GPU"""
+    global runtime, load_error
+    
+    # مرحله ۱: نصب رانتایم
+    if not setup_runtime():
+        load_error = "Failed to setup runtime. Check git clone logs."
+        return
+    
+    # مرحله ۲: دانلود چک‌پوینت‌ها
+    if not download_checkpoints():
+        load_error = "Failed to download checkpoints."
+        return
+    
+    # مرحله ۳: بارگذاری مدل
     try:
-        # تلاش برای بارگذاری ماژول اصلی JoyAI
-        from xvideo.inductor_autotune_fix import install as _install_autotune_fix
+        print("🔄 Loading model into GPU memory...")
+        from xvideo.inductor_autotune_fix import install as _fix
         from xvideo.serving.joyomni_streaming import JoyOmniRuntime
         
-        _install_autotune_fix()
+        _fix()
+        ckpt = Path(os.environ["JOYOMNI_CKPT_ROOT"])
         
-        ckpt_root = Path(os.environ["JOYOMNI_CKPT_ROOT"])
-        rt = JoyOmniRuntime.load(
-            ckpt_root / "dit" / "joyai_video_edit_dit_0811.pth",
-            vae_ckpt=ckpt_root / "vae",
-            text_encoder_ckpt=ckpt_root / "text_encoder",
+        runtime = JoyOmniRuntime.load(
+            ckpt / "dit" / "joyai_video_edit_dit_0811.pth",
+            vae_ckpt=ckpt / "vae",
+            text_encoder_ckpt=ckpt / "text_encoder",
             device="cuda",
             seed=42,
             warmup_height=480,
             warmup_width=840,
         )
-        return rt, None
+        load_error = None
+        print("✅ Model loaded successfully! Ready for inference.")
+        
     except ImportError as e:
-        return None, f"xvideo module not found: {e}. Please ensure JoyAI runtime is installed."
+        load_error = f"xvideo module import failed: {e}\n" \
+                     f"sys.path includes: {[p for p in sys.path if 'JoyAI' in p]}"
+        print(f"❌ {load_error}")
     except Exception as e:
-        return None, f"Model loading failed: {str(e)}"
-
-# بارگذاری اولیه مدل
-print("🔄 Loading JoyAI-Video-Edit model...")
-runtime, error_msg = load_model()
-if error_msg:
-    print(f"⚠️ Warning: {error_msg}")
-else:
-    print("✅ Model loaded successfully!")
+        load_error = f"Model loading error: {str(e)}"
+        print(f"❌ {load_error}")
 
-# ✅ رابط کاربری Gradio (الزامی برای ZeroGPU)
-import gradio as gr
 
+# ============================================================
+# ✅ ۵. تابع ویرایش ویدیو (با GPU decorator)
+# ============================================================
+@spaces.GPU(duration=60, gpu_size="xlarge")
 def edit_video(video_path, prompt):
-    """تابع ویرایش ویدیو"""
+    """ویرایش ویدیو با دستور متنی"""
     if runtime is None:
-        return None, f"❌ Model not available: {error_msg}"
-    
+        return None, f"❌ Model not available:\n{load_error}"
     if not video_path:
         return None, "❌ Please upload a video first."
-    
-    if not prompt.strip():
+    if not prompt or not prompt.strip():
         return None, "❌ Please enter an editing instruction."
     
     try:
-        # TODO: پیاده‌سازی منطق ویرایش با استفاده از runtime
-        # فعلاً یک پاسخ placeholder برمی‌گردانیم
-        return video_path, f"✅ Received: '{prompt}'\n(Runtime ready for streaming inference)"
+        # TODO: جایگزینی با استریمینگ واقعی JoyAI
+        # مثال استفاده از runtime:
+        # output_frames = runtime.edit(video_path, prompt)
+        # output_path = save_frames(output_frames)
+        # return output_path, "✅ Done!"
+        
+        return video_path, (
+            f"✅ Instruction received: '{prompt}'\n"
+            f"🟢 Model is loaded and ready.\n"
+            f"⚠️ Streaming inference pipeline not yet wired.\n"
+            f"💡 See TODO in app.py line ~150 to connect runtime.edit()"
+        )
     except Exception as e:
         return None, f"❌ Editing failed: {str(e)}"
 
-with gr.Blocks(title="JoyAI Video Edit") as demo:
-    gr.Markdown("# 🎬 JoyAI-Video-Edit\nReal-Time Open-Ended Video Editing")
+
+# ============================================================
+# ✅ ۶. رابط کاربری Gradio
+# ============================================================
+import gradio as gr
+
+# بارگذاری اولیه (قبل از ساخت UI)
+print("\n" + "="*60)
+print("🚀 JoyAI-Video-Edit — Initializing...")
+print("="*60)
+load_model()
+
+with gr.Blocks(
+    title="JoyAI Video Edit",
+    theme=gr.themes.Soft(),
+) as demo:
+    gr.Markdown("""
+    # 🎬 JoyAI-Video-Edit
+    ### Real-Time Open-Ended Video Editing with Autoregressive Diffusion (16B)
+    """)
     
     with gr.Row():
-        with gr.Column():
-            input_video = gr.Video(label="Input Video", sources=["upload", "webcam"])
+        with gr.Column(scale=1):
+            input_video = gr.Video(
+                label="📹 Input Video",
+                sources=["upload", "webcam"],
+                height=360,
+            )
             edit_prompt = gr.Textbox(
-                label="Editing Instruction",
-                placeholder="e.g., Transform into watercolor style, Add sunglasses...",
-                lines=2
+                label="✏️ Editing Instruction",
+                placeholder="Examples:\n• Transform into watercolor style\n• Add sunglasses to the person\n• Replace background with beach\n• Remove the cat on the left",
+                lines=3,
+            )
+            submit_btn = gr.Button(
+                "🚀 Edit Video",
+                variant="primary",
+                size="lg",
             )
-            submit_btn = gr.Button("✏️ Edit Video", variant="primary")
         
-        with gr.Column():
-            output_video = gr.Video(label="Edited Output")
-            status_log = gr.Textbox(label="Status", interactive=False)
+        with gr.Column(scale=1):
+            output_video = gr.Video(label="🎞️ Edited Output", height=360)
+            status_log = gr.Textbox(
+                label="📋 Status",
+                interactive=False,
+                lines=6,
+                value=f"{'✅ Model Ready' if runtime else '⚠️ ' + (load_error or 'Loading...')}",
+            )
     
     submit_btn.click(
         fn=edit_video,
         inputs=[input_video, edit_prompt],
-        outputs=[output_video, status_log]
+        outputs=[output_video, status_log],
     )
     
     gr.Markdown("""
-    ### ⚠️ Notes
-    - This space requires **ZeroGPU (xlarge tier)**
-    - First run may take 1-2 minutes for model warmup
-    - Supports: style transfer, object edit, subject replacement, background change
+    ---
+    ### ℹ️ About
+    | Item | Detail |
+    |------|--------|
+    | **Model** | JoyAI-Video-Edit (16B Autoregressive Diffusion) |
+    | **Hardware** | ZeroGPU xlarge (RTX Pro 6000 Blackwell, 96GB) |
+    | **Capabilities** | Style transfer, object edit, subject add/remove/replace, background change |
+    | **Paper** | [arXiv:2608.03974](https://arxiv.org/abs/2608.03974) |
+    | **Source** | [jd-opensource/JoyAI-Video-Edit](https://github.com/jd-opensource/JoyAI-Video-Edit) |
+    
+    ⏱️ **First run:** Model download (~30GB) + warmup may take 15-25 minutes.
     """)
 
-# ✅ جلوگیری از خروج فوری سرور
+
+# ============================================================
+# ✅ ۷. راه‌اندازی سرور
+# ============================================================
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        ssr_mode=False,
+        show_api=False,
+    )