Update app.py

app.py

@@ -16,23 +16,23 @@ import subprocess
 import os
 import sys
 
-# --- Setup: Clone repository and add it to Python Path ---
-# This section ensures all necessary code and model files are available.
+# --- Setup: Clone repository and Change Working Directory ---
+# This is the most robust way to ensure all relative paths work correctly.
 
 # 1. Clone the repository with all its files
 subprocess.run("git lfs install", shell=True, check=True)
-if not os.path.exists("SeedVR2-3B"):
-    print("Cloning SeedVR2-3B repository...")
-    subprocess.run("git clone https://huggingface.co/spaces/ByteDance-Seed/SeedVR2-3B", shell=True, check=True)
+repo_dir_name = "SeedVR2-3B"
+if not os.path.exists(repo_dir_name):
+    print(f"Cloning {repo_dir_name} repository...")
+    subprocess.run(f"git clone https://huggingface.co/spaces/ByteDance-Seed/{repo_dir_name}", shell=True, check=True)
 
-# 2. Add the cloned repository's directory to Python's module search path
-repo_dir = "SeedVR2-3B"
-# This allows us to import modules like 'data', 'common', etc., from the cloned repo.
-sys.path.insert(0, os.path.abspath(repo_dir))
-print(f"Repository directory '{os.path.abspath(repo_dir)}' added to Python path.")
+# 2. Change the current working directory to the repository's root
+# CORREÇÃO PRINCIPAL: Isso resolve todos os problemas de caminho relativo.
+os.chdir(repo_dir_name)
+print(f"Changed working directory to: {os.getcwd()}")
 
 # --- Main Application Code ---
-# All file paths will now be relative to the cloned repository directory.
+# Now that we are inside the repo, all imports and file loads will work naturally.
 
 import torch
 import mediapy
@@ -51,7 +51,7 @@ import torchvision.transforms as T
 from torchvision.transforms import Compose, Lambda, Normalize
 from torchvision.io.video import read_video
 
-# Imports from the cloned repository
+# Imports from the repository (will now work directly)
 from data.image.transforms.divisible_crop import DivisibleCrop
 from data.image.transforms.na_resize import NaResize
 from data.video.transforms.rearrange import Rearrange
@@ -63,9 +63,8 @@ from common.partition import partition_by_size
 from projects.video_diffusion_sr.infer import VideoDiffusionInfer
 from common.distributed.ops import sync_data
 
-# Check for color_fix utility
-color_fix_path = os.path.join(repo_dir, "projects/video_diffusion_sr/color_fix.py")
-if os.path.exists(color_fix_path):
+# Check for color_fix utility (using relative path)
+if os.path.exists("projects/video_diffusion_sr/color_fix.py"):
     from projects.video_diffusion_sr.color_fix import wavelet_reconstruction
     use_colorfix = True
 else:
@@ -78,7 +77,7 @@ os.environ["MASTER_PORT"] = "12355"
 os.environ["RANK"] = str(0)
 os.environ["WORLD_SIZE"] = str(1)
 
-# CORREÇÃO: Usar sys.executable para chamar o pip corretamente
+# Use sys.executable to ensure we use the correct pip
 python_executable = sys.executable
 subprocess.run(
     [python_executable, "-m", "pip", "install", "flash-attn", "--no-build-isolation"],
@@ -86,9 +85,8 @@ subprocess.run(
     check=True
 )
 
-apex_wheel_path = os.path.join(repo_dir, "apex-0.1-cp310-cp310-linux_x86_64.whl")
+apex_wheel_path = "apex-0.1-cp310-cp310-linux_x86_64.whl"
 if os.path.exists(apex_wheel_path):
-    # CORREÇÃO: Usar sys.executable aqui também
     subprocess.run([python_executable, "-m", "pip", "install", apex_wheel_path], check=True)
     print("✅ Apex setup completed.")
 
@@ -99,10 +97,11 @@ def configure_sequence_parallel(sp_size):
         init_sequence_parallel(sp_size)
 
 def configure_runner(sp_size):
-    config_path = os.path.join(repo_dir, 'configs_3b', 'main.yaml')
-    checkpoint_path = os.path.join(repo_dir, 'ckpts', 'seedvr2_ema_3b.pth')
+    # Paths are now simple and relative to the repo root
+    config_path = 'configs_3b/main.yaml'
+    checkpoint_path = 'ckpts/seedvr2_ema_3b.pth'
     
-    config = load_config(config_path)
+    config = load_config(config_path) # This will now work correctly
     runner = VideoDiffusionInfer(config)
     OmegaConf.set_readonly(runner.config, False)
     
@@ -158,8 +157,9 @@ def generation_loop(video_path, seed=666, fps_out=24, batch_size=1, cfg_scale=1.
     def _extract_text_embeds():
         positive_prompts_embeds = []
         for _ in original_videos_local:
-            text_pos_embeds = torch.load(os.path.join(repo_dir, 'pos_emb.pt'))
-            text_neg_embeds = torch.load(os.path.join(repo_dir, 'neg_emb.pt'))
+            # Paths are now simple
+            text_pos_embeds = torch.load('pos_emb.pt')
+            text_neg_embeds = torch.load('neg_emb.pt')
             positive_prompts_embeds.append({"texts_pos": [text_pos_embeds], "texts_neg": [text_neg_embeds]})
         gc.collect()
         torch.cuda.empty_cache()
@@ -276,13 +276,12 @@ def generation_loop(video_path, seed=666, fps_out=24, batch_size=1, cfg_scale=1.
 # --- Gradio UI ---
 
 with gr.Blocks(title="SeedVR2: One-Step Video Restoration") as demo:
-    logo_path = os.path.join(repo_dir, "assets/seedvr_logo.png")
+    # Use an absolute path for the Gradio file source to be safe
+    logo_path = os.path.abspath("assets/seedvr_logo.png")
     gr.HTML(f"""
-        <div style='text-align:center; margin-bottom: 10px;'>
-            <img src='file/{logo_path}' style='height:40px;' alt='SeedVR logo'/>
-        </div>
-        <p><b>Official Gradio demo</b> for <a href='https://github.com/ByteDance-Seed/SeedVR' target='_blank'><b>SeedVR2: One-Step Video Restoration via Diffusion Adversarial Post-Training</b></a>.<br>
-        🔥 <b>SeedVR2</b> is a one-step image and video restoration algorithm for real-world and AIGC content.</p>
+        
+        <a href='https://github.com/ByteDance-Seed/SeedVR' target='_blank'><b>SeedVR2: One-Step Video Restoration via Diffusion Adversarial Post-Training</b></a>.<br>
+        🔥 <b>SeedVR2</b> is a one-step image and video restoration algorithm for real-world and AIGC content.
     """)
 
     with gr.Row():