Update app.py

app.py

@@ -11,23 +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.
@@ -35,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 ---
@@ -47,14 +33,38 @@ 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.
+# ORDEM CORRETA: Instalar requisitos, DEPOIS compilar apex/flash-attn, DEPOIS baixar modelos.
+
+python_executable = sys.executable
+
+# **CORREÇÃO CRÍTICA: 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 instaladas.")
+
+
+# **Compilar dependências otimizadas para a GPU L40S (Ada Lovelace)**
+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)
+
+# Instala o Apex a partir da fonte clonada, o que força a compilação para a GPU L40S
+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.")
 
+
+# **Download dos modelos e dados**
 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()
@@ -62,15 +72,13 @@ def load_file_from_url(url, model_dir=None, progress=True, file_name=None):
     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
+    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/SeedVR2-3B/resolve/main/ema_vae.pth',
 	'dit': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/seedvr2_ema_3b.pth',
@@ -78,9 +86,7 @@ pretrain_model_url = {
     'neg_emb': 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/neg_emb.pt',
 }
 
-# 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 '.'
@@ -88,29 +94,10 @@ for key, url in pretrain_model_url.items():
     if not os.path.exists(target_path):
         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')
-
-# --- 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.")
+print("✅ Modelos e dados de exemplo baixados.")
 
 # --- ETAPA 4: Execução do Código Principal da Aplicação ---
 import mediapy
@@ -142,7 +129,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"):
@@ -186,7 +172,7 @@ def generation_step(runner, text_embeds_dict, cond_latents):
         video_tensors = runner.inference(noises=noises, conditions=conditions, dit_offload=False, **text_embeds_dict)
     return [rearrange(video, "c t h w -> t c h w") for video in video_tensors]
 
-@spaces.GPU
+
 def generation_loop(video_path, seed=666, fps_out=24, batch_size=1, cfg_scale=1.0, cfg_rescale=0.0, sample_steps=1, res_h=1280, res_w=720, sp_size=1):
     if video_path is None: return None, None, None
     runner = configure_runner(1)
@@ -228,7 +214,9 @@ 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"""
-        
+        <div style='text-align:center; margin-bottom: 10px;'>
+            <img src='file/{os.path.abspath("assets/seedvr_logo.png")}' style='height:40px;' alt='SeedVR logo'/>
+        </div>
         <p><b>Demonstração oficial do Gradio</b> para 
         <a href='https://github.com/ByteDance-Seed/SeedVR' target='_blank'>
         <b>SeedVR2: One-Step Video Restoration via Diffusion Adversarial Post-Training</b></a>.<br>