Update app.py

app.py

@@ -15,28 +15,6 @@ import spaces
 import subprocess
 import os
 import sys
-
-# --- Setup: Clone repository, Change Directory, and Update Python Path ---
-# Esta é a abordagem definitiva para corrigir todos os problemas de caminho.
-
-# 1. Clone o repositório
-repo_dir_name = "SeedVR2-3B"
-if not os.path.exists(repo_dir_name):
-    print(f"Clonando o repositório {repo_dir_name}...")
-    subprocess.run(f"git clone https://huggingface.co/spaces/ByteDance-Seed/{repo_dir_name}", shell=True, check=True)
-
-# 2. Mude o diretório de trabalho atual para a raiz do repositório.
-os.chdir(repo_dir_name)
-print(f"Diretório de trabalho alterado para: {os.getcwd()}")
-
-# 3. Adicione explicitamente o novo diretório de trabalho ao caminho do sistema do Python.
-sys.path.insert(0, os.path.abspath('.'))
-print(f"Diretório atual adicionado ao sys.path: {os.path.abspath('.')}")
-
-
-# --- Código Principal da Aplicação ---
-# Agora, todas as importações e cargas de arquivos devem funcionar corretamente.
-
 import torch
 import mediapy
 from einops import rearrange
@@ -47,6 +25,8 @@ import gc
 from PIL import Image
 import gradio as gr
 from pathlib import Path
+from urllib.parse import urlparse
+from torch.hub import download_url_to_file, get_dir
 import shlex
 import uuid
 import mimetypes
@@ -54,33 +34,69 @@ import torchvision.transforms as T
 from torchvision.transforms import Compose, Lambda, Normalize
 from torchvision.io.video import read_video
 
-# Importações do repositório (agora funcionarão)
-from data.image.transforms.divisible_crop import DivisibleCrop
-from data.image.transforms.na_resize import NaResize
-from data.video.transforms.rearrange import Rearrange
-from common.config import load_config
-from common.distributed import init_torch
-from common.distributed.advanced import init_sequence_parallel
-from common.seed import set_seed
-from common.partition import partition_by_size
-from projects.video_diffusion_sr.infer import VideoDiffusionInfer
-from common.distributed.ops import sync_data
+# --- Lógica de Download de Arquivos (do script original) ---
+
+def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
+    """Carrega um arquivo de um URL http, baixando modelos se necessário."""
+    if model_dir is None:
+        hub_dir = get_dir()
+        model_dir = os.path.join(hub_dir, 'checkpoints')
+
+    os.makedirs(model_dir, exist_ok=True)
+
+    parts = urlparse(url)
+    filename = os.path.basename(parts.path)
+    if file_name is not None:
+        filename = file_name
+    cached_file = os.path.abspath(os.path.join(model_dir, filename))
+    if not os.path.exists(cached_file):
+        print(f'Baixando: "{url}" para {cached_file}\n')
+        download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
+    return cached_file
+
+ckpt_dir = Path('./ckpts')
+if not ckpt_dir.exists():
+	ckpt_dir.mkdir()
+
+pretrain_model_url = {
+	'vae': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/ema_vae.pth',
+	'dit': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/seedvr2_ema_3b.pth',
+    'pos_emb': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/pos_emb.pt',
+    'neg_emb': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/neg_emb.pt',
+    'apex': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/apex-0.1-cp310-cp310-linux_x86_64.whl'
+}
+
+# Baixa os pesos e dependências se não existirem
+if not os.path.exists('./ckpts/seedvr2_ema_3b.pth'):
+	load_file_from_url(url=pretrain_model_url['dit'], model_dir='./ckpts/', progress=True)
+if not os.path.exists('./ckpts/ema_vae.pth'):
+	load_file_from_url(url=pretrain_model_url['vae'], model_dir='./ckpts/', progress=True)
+if not os.path.exists('./pos_emb.pt'):
+	load_file_from_url(url=pretrain_model_url['pos_emb'], model_dir='./', progress=True)
+if not os.path.exists('./neg_emb.pt'):
+	load_file_from_url(url=pretrain_model_url['neg_emb'], model_dir='./', progress=True)
+if not os.path.exists('./apex-0.1-cp310-cp310-linux_x86_64.whl'):
+	load_file_from_url(url=pretrain_model_url['apex'], model_dir='./', progress=True)
+
+# Baixa os vídeos de exemplo
+torch.hub.download_url_to_file(
+	'https://huggingface.co/datasets/Iceclear/SeedVR_VideoDemos/resolve/main/seedvr_videos_crf23/aigc1k/23_1_lq.mp4',
+	'01.mp4')
+torch.hub.download_url_to_file(
+	'https://huggingface.co/datasets/Iceclear/SeedVR_VideoDemos/resolve/main/seedvr_videos_crf23/aigc1k/28_1_lq.mp4',
+	'02.mp4')
+torch.hub.download_url_to_file(
+	'https://huggingface.co/datasets/Iceclear/SeedVR_VideoDemos/resolve/main/seedvr_videos_crf23/aigc1k/2_1_lq.mp4',
+	'03.mp4')
 
-# Verifica o utilitário color_fix (usando caminho relativo)
-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:
-    use_colorfix = False
-    print('Atenção!!!!!! A correção de cor não está disponível!')
 
 # --- Configuração de Ambiente e Dependências ---
+
 os.environ["MASTER_ADDR"] = "127.0.0.1"
 os.environ["MASTER_PORT"] = "12355"
 os.environ["RANK"] = str(0)
 os.environ["WORLD_SIZE"] = str(1)
 
-# Use sys.executable para garantir que estamos usando o pip correto
 python_executable = sys.executable
 subprocess.run(
     [python_executable, "-m", "pip", "install", "flash-attn", "--no-build-isolation"],
@@ -91,25 +107,40 @@ subprocess.run(
 apex_wheel_path = "apex-0.1-cp310-cp310-linux_x86_64.whl"
 if os.path.exists(apex_wheel_path):
     print("Instalando o Apex a partir do arquivo wheel...")
-    # CORREÇÃO: Usar --force-reinstall e --no-cache-dir para garantir uma instalação limpa.
     subprocess.run(
         [python_executable, "-m", "pip", "install", "--force-reinstall", "--no-cache-dir", apex_wheel_path], 
         check=True
     )
     print("✅ Configuração do Apex concluída.")
 
-# --- Funções Principais ---
+# --- Código Principal da Aplicação ---
+
+from data.image.transforms.divisible_crop import DivisibleCrop
+from data.image.transforms.na_resize import NaResize
+from data.video.transforms.rearrange import Rearrange
+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:
+    use_colorfix = False
+    print('Atenção!!!!!! A correção de cor não está disponível!')
+from common.config import load_config
+from common.distributed import init_torch
+from common.distributed.advanced import init_sequence_parallel
+from common.seed import set_seed
+from common.partition import partition_by_size
+from projects.video_diffusion_sr.infer import VideoDiffusionInfer
+from common.distributed.ops import sync_data
 
 def configure_sequence_parallel(sp_size):
     if sp_size > 1:
         init_sequence_parallel(sp_size)
 
 def configure_runner(sp_size):
-    # Os caminhos agora são simples e relativos à raiz do repositório
     config_path = 'configs_3b/main.yaml'
     checkpoint_path = 'ckpts/seedvr2_ema_3b.pth'
     
-    config = load_config(config_path) # Isto agora funcionará corretamente
+    config = load_config(config_path)
     runner = VideoDiffusionInfer(config)
     OmegaConf.set_readonly(runner.config, False)
     
@@ -165,7 +196,6 @@ 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:
-            # Os caminhos agora são simples
             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]})
@@ -284,8 +314,7 @@ def generation_loop(video_path, seed=666, fps_out=24, batch_size=1, cfg_scale=1.
 # --- UI do Gradio ---
 
 with gr.Blocks(title="SeedVR2: Restauração de Vídeo em Um Passo") as demo:
-    # Use um caminho absoluto para o arquivo de logo do Gradio para segurança
-    logo_path = os.path.abspath("assets/seedvr_logo.png")
+    logo_path = "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'/>
@@ -310,6 +339,16 @@ with gr.Blocks(title="SeedVR2: Restauração de Vídeo em Um Passo") as demo:
 
     run_button.click(fn=generation_loop, inputs=[input_file, seed, fps], outputs=[output_image, output_video, download_link])
 
+    # Seção de Exemplos, que agora funcionará pois os vídeos são baixados
+    gr.Examples(
+        examples=[
+            ["./01.mp4", 4, 24],
+            ["./02.mp4", 4, 24],
+            ["./03.mp4", 4, 24],
+        ],
+        inputs=[input_file, seed, fps]
+    )
+
     gr.HTML("""
         <hr>
         <p>Se você achou o SeedVR útil, por favor ⭐ o <a href='https://github.com/ByteDance-Seed/SeedVR' target='_blank'>repositório no GitHub</a>: