Update app.py

app.py

@@ -11,22 +11,10 @@
 # // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # // See the License for the specific language governing permissions and
 # // limitations under the License.
-
-import torch
-import torch.distributed as dist
+import spaces
+import subprocess
 import os
-import gc
-import logging
 import sys
-import subprocess
-from pathlib import Path
-from urllib.parse import urlparse
-from torch.hub import download_url_to_file
-import gradio as gr
-import mediapy
-from einops import rearrange
-import shutil
-from omegaconf import OmegaConf
 
 # --- ETAPA 1: Preparação do Ambiente ---
 # Clonar o repositório para garantir que todas as pastas de código (data, common, etc.) existam.
@@ -34,7 +22,6 @@ from omegaconf import OmegaConf
 repo_dir_name = "SeedVR2-3B"
 if not os.path.exists(repo_dir_name):
     print(f"Clonando o repositório {repo_dir_name} para obter todo o código-fonte...")
-    # Usamos --depth 1 para um clone mais rápido, já que não precisamos do histórico
     subprocess.run(f"git clone --depth 1 https://huggingface.co/spaces/ByteDance-Seed/{repo_dir_name}", shell=True, check=True)
 
 # --- ETAPA 2: Configuração dos Caminhos ---
@@ -45,31 +32,62 @@ print(f"Diretório de trabalho alterado para: {os.getcwd()}")
 sys.path.insert(0, os.path.abspath('.'))
 print(f"Diretório atual adicionado ao sys.path para importações.")
 
-# --- ETAPA 3: Instalação de Dependências e Download de Modelos ---
-# Agora que estamos no diretório correto, podemos prosseguir.
+# --- ETAPA 3: Instalação de Dependências (NA ORDEM CORRETA E COM A FLAG CORRETA) ---
+
+python_executable = sys.executable
+
+# **PASSO 3.1: Instalar requisitos PRIMEIRO para ter o PyTorch disponível**
+print("Instalando dependências a partir do requirements.txt (isso inclui o PyTorch)...")
+subprocess.run([python_executable, "-m", "pip", "install", "-r", "requirements.txt"], check=True)
+print("✅ Dependências básicas (incluindo PyTorch) instaladas.")
+
+
+# **PASSO 3.2: Compilar dependências otimizadas para a GPU**
+print("Instalando flash-attn compilando do zero...")
+subprocess.run([python_executable, "-m", "pip", "install", "--force-reinstall", "--no-cache-dir", "flash-attn"], check=True)
+
+print("Clonando e compilando o Apex do zero (isso pode demorar um pouco)...")
+if not os.path.exists("apex"):
+    subprocess.run("git clone https://github.com/NVIDIA/apex", shell=True, check=True)
 
+# **CORREÇÃO FINAL: Adicionar a flag --no-build-isolation**
+# Isso força o build a usar o ambiente atual (onde o torch já foi instalado)
+# em vez de criar um ambiente isolado e vazio.
+print("Compilando e instalando o Apex...")
+subprocess.run(
+    [
+        python_executable, "-m", "pip", "install",
+        "--no-build-isolation",  # A FLAG CRÍTICA QUE RESOLVE O PROBLEMA
+        "-v",
+        "--disable-pip-version-check",
+        "--no-cache-dir",
+        "--global-option=--cpp_ext",
+        "--global-option=--cuda_ext",
+        "./apex"
+    ],
+    check=True
+)
+print("✅ Configuração do Apex concluída.")
+
+
+# **PASSO 3.3: Download dos modelos e dados de exemplo**
 import torch
 from pathlib import Path
 from urllib.parse import urlparse
 from torch.hub import download_url_to_file, get_dir
 
-# Função de download do original
 def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
     if model_dir is None:
-        hub_dir = get_dir()
-        model_dir = os.path.join(hub_dir, 'checkpoints')
+        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
+    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
 
-# URLs dos modelos
 pretrain_model_url = {
 	'vae': 'https://huggingface.co/ByteDance-Seed/SeedVR-7B/resolve/main/ema_vae.pth',
 	'dit': 'https://huggingface.co/ByteDance-Seed/SeedVR-7B/resolve/main/seedvr_ema_7b.pth',
@@ -78,40 +96,18 @@ pretrain_model_url = {
     'apex': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/apex-0.1-cp310-cp310-linux_x86_64.whl'
 }
 
-
-# Criar diretório de checkpoints e baixar modelos
-ckpt_dir = Path('./ckpts')
-ckpt_dir.mkdir(exist_ok=True)
+ckpt_dir = Path('./ckpts'); ckpt_dir.mkdir(exist_ok=True)
 for key, url in pretrain_model_url.items():
     filename = os.path.basename(url)
     model_dir = './ckpts' if key in ['vae', 'dit'] else '.'
-    target_path = os.path.join(model_dir, filename)
-    if not os.path.exists(target_path):
+    if not os.path.exists(os.path.join(model_dir, filename)):
         load_file_from_url(url=url, model_dir=model_dir, progress=True, file_name=filename)
 
-# Baixar 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')
+print("✅ Modelos e dados de exemplo baixados.")
 
-# --- REFINAMENTO: Compilar dependências do zero para a GPU L40S (Ada Lovelace) ---
-python_executable = sys.executable
-
-print("Instalando flash-attn compilando do zero...")
-# Força a reinstalação a partir do zero para garantir que seja compilado para a GPU atual
-subprocess.run([python_executable, "-m", "pip", "install", "--force-reinstall", "--no-cache-dir", "flash-attn"], check=True)
-
-print("Clonando e compilando o Apex do zero...")
-if not os.path.exists("apex"):
-    subprocess.run("git clone https://github.com/NVIDIA/apex", shell=True, check=True)
-
-# Instala o Apex a partir da fonte clonada, o que força a compilação para a GPU L40S
-# As flags --cpp_ext e --cuda_ext são essenciais para a compilação
-subprocess.run(
-    [python_executable, "-m", "pip", "install", "-v", "--disable-pip-version-check", "--no-cache-dir", "--global-option=--cpp_ext", "--global-option=--cuda_ext", "./apex"],
-    check=True
-)
-print("✅ Configuração do Apex concluída.")
 
 # --- ETAPA 4: Execução do Código Principal da Aplicação ---
 import mediapy
@@ -143,7 +139,6 @@ os.environ["MASTER_ADDR"] = "127.0.0.1"
 os.environ["MASTER_PORT"] = "12355"
 os.environ["RANK"] = str(0)
 os.environ["WORLD_SIZE"] = str(1)
-# Adiciona uma variável de ambiente que pode ajudar o PyTorch a debugar erros de CUDA
 os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
 
 if os.path.exists("projects/video_diffusion_sr/color_fix.py"):
@@ -229,10 +224,8 @@ def generation_loop(video_path, seed=666, fps_out=24, batch_size=1, cfg_scale=1.
 
 with gr.Blocks(title="SeedVR2: Restauração de Vídeo em Um Passo") as demo:
     gr.HTML(f"""
-        
-        <p>
         <a href='https://github.com/ByteDance-Seed/SeedVR' target='_blank'>
-        <b>SeedVR2 7B: One-Step Video Restoration via Diffusion Adversarial Post-Training</b></a>.<br>
+        <b>SeedVR2 7b: One-Step Video Restoration via Diffusion Adversarial Post-Training</b></a>.<br>
         🔥 <b>SeedVR2</b> é um algoritmo de restauração de imagem e vídeo em um passo para conteúdo do mundo real e AIGC.
         </p>
     """)
@@ -255,17 +248,5 @@ with gr.Blocks(title="SeedVR2: Restauração de Vídeo em Um Passo") as demo:
         ],
         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>:</p>
-        <a href="https://github.com/ByteDance-Seed/SeedVR" target="_blank">
-            <img src="https://img.shields.io/github/stars/ByteDance-Seed/SeedVR?style=social" alt="GitHub Stars">
-        </a>
-        <h4>Aviso</h4>
-        <p>Esta demonstração suporta até <b>720p e 121 frames para vídeos ou imagens 2k</b>.  
-        Para outros casos de uso, verifique o <a href='https://github.com/ByteDance-Seed/SeedVR' target='_blank'>repositório no GitHub</a>.</p>
-        <h4>Limitações</h4>
-        <p>Pode falhar em degradações pesadas ou em clipes AIGC com pouco movimento, causando excesso de nitidez ou restauração inadequada.</p>
-    """)
+    
 demo.queue().launch(share=True)