Update app.py

app.py

@@ -2,7 +2,7 @@
 # //
 # // Licensed under the Apache License, Version 2.0 (the "License");
 # // you may not use this file except in compliance with the License.
-# // You may not obtain a copy of the License at
+# // You may obtain a copy of the License at
 # //
 # //     http://www.apache.org/licenses/LICENSE-2.0
 # //
@@ -11,23 +11,12 @@
 # // 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.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: Clonar o Repositório Oficial do GitHub ---
+# --- ETAPA 1: Clonar o Repositório do GitHub ---
 repo_name = "SeedVR"
 if not os.path.exists(repo_name):
     print(f"Clonando o repositório {repo_name} do GitHub...")
@@ -37,14 +26,25 @@ if not os.path.exists(repo_name):
 os.chdir(repo_name)
 print(f"Diretório de trabalho alterado para: {os.getcwd()}")
 
-# Adicionar o diretório ao path do Python para que as importações funcionem
 sys.path.insert(0, os.path.abspath('.'))
 print(f"Diretório atual adicionado ao sys.path.")
 
-# --- ETAPA 3: Instalar Dependências Conforme as Instruções ---
+# --- ETAPA 3: Instalar Dependências Corretamente ---
 python_executable = sys.executable
-print("Instalando dependências do requirements.txt...")
-subprocess.run([python_executable, "-m", "pip", "install", "-r", "requirements.txt"], check=True)
+
+# CORREÇÃO: Forçar uma versão do NumPy < 2.0 para evitar conflitos de compatibilidade.
+print("Instalando NumPy compatível...")
+subprocess.run([python_executable, "-m", "pip", "install", "numpy<2.0"], check=True)
+
+# Filtrar requirements.txt para evitar conflitos com torch/torchvision pré-instalados
+print("Filtrando requirements.txt...")
+with open("requirements.txt", "r") as f_in, open("filtered_requirements.txt", "w") as f_out:
+    for line in f_in:
+        if not line.strip().startswith(('torch', 'torchvision')):
+            f_out.write(line)
+
+print("Instalando dependências filtradas...")
+subprocess.run([python_executable, "-m", "pip", "install", "-r", "filtered_requirements.txt"], check=True)
 
 print("Instalando flash-attn...")
 subprocess.run([python_executable, "-m", "pip", "install", "flash-attn==2.5.9.post1", "--no-build-isolation"], check=True)
@@ -53,7 +53,6 @@ from pathlib import Path
 from urllib.parse import urlparse
 from torch.hub import download_url_to_file, get_dir
 
-# Função auxiliar para downloads
 def load_file_from_url(url, model_dir='.', progress=True, file_name=None):
     os.makedirs(model_dir, exist_ok=True)
     if not file_name:
@@ -65,7 +64,6 @@ def load_file_from_url(url, model_dir='.', progress=True, file_name=None):
         download_url_to_file(url, cached_file, hash_prefix=None, progress=progress)
     return cached_file
 
-# Baixar e instalar Apex pré-compilado (crucial para o ambiente do Spaces)
 apex_url = 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/apex-0.1-cp310-cp310-linux_x86_64.whl'
 apex_wheel_path = load_file_from_url(url=apex_url)
 print("Instalando Apex a partir do wheel baixado...")
@@ -74,6 +72,8 @@ print("✅ Configuração do Apex concluída.")
 
 # --- ETAPA 4: Baixar os Modelos Pré-treinados ---
 print("Baixando modelos pré-treinados...")
+import torch
+
 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',
@@ -89,7 +89,6 @@ for key, url in pretrain_model_url.items():
 
 
 # --- ETAPA 5: Executar a Aplicação Principal ---
-import torch
 import mediapy
 from einops import rearrange
 from omegaconf import OmegaConf
@@ -109,9 +108,7 @@ 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
 
@@ -120,11 +117,9 @@ os.environ["MASTER_PORT"] = "12355"
 os.environ["RANK"] = str(0)
 os.environ["WORLD_SIZE"] = str(1)
 
-if os.path.exists("projects/video_diffusion_sr/color_fix.py"):
+use_colorfix = os.path.exists("projects/video_diffusion_sr/color_fix.py")
+if use_colorfix:
     from projects.video_diffusion_sr.color_fix import wavelet_reconstruction
-    use_colorfix = True
-else:
-    use_colorfix = False
 
 def configure_runner():
     config = load_config('configs_7b/main.yaml')
@@ -139,10 +134,9 @@ def configure_runner():
 
 def generation_step(runner, text_embeds_dict, cond_latents):
     def _move_to_cuda(x): return [i.to("cuda") for i in x]
-    noises = [torch.randn_like(latent) for latent in cond_latents]
-    aug_noises = [torch.randn_like(latent) for latent in cond_latents]
+    noises, aug_noises = [torch.randn_like(l) for l in cond_latents], [torch.randn_like(l) for l in cond_latents]
     noises, aug_noises, cond_latents = sync_data((noises, aug_noises, cond_latents), 0)
-    noises, aug_noises, cond_latents = list(map(_move_to_cuda, (noises, aug_noises, cond_latents)))
+    noises, aug_noises, cond_latents = map(_move_to_cuda, (noises, aug_noises, cond_latents))
     def _add_noise(x, aug_noise):
         t = torch.tensor([100.0], device="cuda")
         shape = torch.tensor(x.shape[1:], device="cuda")[None]
@@ -157,13 +151,28 @@ def generation_step(runner, text_embeds_dict, cond_latents):
 def generation_loop(video_path, seed=666, fps_out=24):
     if video_path is None: return None, None, None
     runner = configure_runner()
-    text_embeds = {"texts_pos": [torch.load('pos_emb.pt').to("cuda")], "texts_neg": [torch.load('neg_emb.pt').to("cuda")]}
+    # Adicionado `weights_only=True` para segurança e para suprimir o aviso
+    text_embeds = {
+        "texts_pos": [torch.load('pos_emb.pt', weights_only=True).to("cuda")], 
+        "texts_neg": [torch.load('neg_emb.pt', weights_only=True).to("cuda")]
+    }
     runner.configure_diffusion()
     set_seed(int(seed))
     os.makedirs("output", exist_ok=True)
-    video_transform = Compose([NaResize(1024), DivisibleCrop(16), Normalize(0.5, 0.5), Rearrange("t c h w -> c t h w")])
+    
+    # CORREÇÃO: Fornecer os argumentos que faltam para NaResize.
+    res_h, res_w = 1280, 720
+    transform = Compose([
+        NaResize(resolution=(res_h * res_w)**0.5, mode="area", downsample_only=False),
+        Lambda(lambda x: torch.clamp(x, 0.0, 1.0)),
+        DivisibleCrop((16, 16)), 
+        Normalize(0.5, 0.5), 
+        Rearrange("t c h w -> c t h w")
+    ])
+
     media_type, _ = mimetypes.guess_type(video_path)
     is_video = media_type and media_type.startswith("video")
+    
     if is_video:
         video, _, _ = read_video(video_path, output_format="TCHW")
         video = video[:121] / 255.0
@@ -171,12 +180,14 @@ def generation_loop(video_path, seed=666, fps_out=24):
     else:
         video = T.ToTensor()(Image.open(video_path).convert("RGB")).unsqueeze(0)
         output_path = os.path.join("output", f"{uuid.uuid4()}.png")
-    cond_latents = [video_transform(video.to("cuda"))]
+
+    cond_latents = [transform(video.to("cuda"))]
     ori_length = cond_latents[0].size(2)
     cond_latents = runner.vae_encode(cond_latents)
     samples = generation_step(runner, text_embeds, cond_latents)
     sample = samples[0][:ori_length].cpu()
     sample = rearrange(sample, "t c h w -> t h w c").clip(-1, 1).add(1).mul(127.5).byte().numpy()
+
     if is_video:
         mediapy.write_video(output_path, sample, fps=fps_out)
         return None, output_path, output_path
@@ -185,7 +196,14 @@ def generation_loop(video_path, seed=666, fps_out=24):
         return output_path, None, output_path
 
 with gr.Blocks(title="SeedVR") 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;'/></div>...""")
+    gr.HTML(f"""
+        
+        <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>
+        🔥 <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>
+    """)
     with gr.Row():
         input_file = gr.File(label="Carregar Imagem ou Vídeo")
         with gr.Column():