Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py (VERSÃO FINAL COMPLETA PARA GRADIO)
+# app.py (VERSÃO FINAL COM MONKEY PATCHING)
 
 import gradio as gr
 import os
@@ -10,108 +10,111 @@ from pathlib import Path
 from huggingface_hub import snapshot_download
 
 # --- BLOCO DE CONFIGURAÇÃO E DOWNLOAD DE MODELO ---
-
-# Diretório de trabalho principal e diretório do código SeedVR
 APP_DIR = "/app"
 SEEDVR_DIR = os.path.join(APP_DIR, "SeedVR")
-
-# Usamos /tmp, um diretório com permissão de escrita garantida, para modelos e arquivos temporários.
 MODEL_CACHE_DIR = "/tmp/models"
 CKPTS_DIR = os.path.join(MODEL_CACHE_DIR, "ckpts")
 
-# Garante que o diretório de checkpoints exista antes de qualquer coisa
 os.makedirs(CKPTS_DIR, exist_ok=True)
 
-# Verifica se um arquivo de modelo chave já existe para evitar redownloads a cada reinício.
 if not Path(CKPTS_DIR).joinpath("seedvr2_ema_3b.pth").exists():
-    print("Baixando os checkpoints do modelo para /tmp/models/ckpts... Isso pode levar alguns minutos.")
+    print("Baixando os checkpoints do modelo para /tmp/models/ckpts...")
     snapshot_download(
         repo_id="ByteDance-Seed/SeedVR2-3B",
         local_dir=CKPTS_DIR,
-        local_dir_use_symlinks=False, # Parâmetro obsoleto, mas mantido por segurança
-        allow_patterns=["*.pth", "*.pt"] # Baixa apenas os arquivos de modelo necessários
+        local_dir_use_symlinks=False,
+        allow_patterns=["*.pth", "*.pt"]
     )
     print("Download do modelo concluído.")
 else:
-    print("Checkpoints do modelo já existem em /tmp. Pulando o download.")
+    print("Checkpoints do modelo já existem em /tmp.")
 
 # --------------------------------------------------------------------
 
 def run_inference(video_path, seed, res_h, res_w):
-    """
-    Função principal que será chamada pela interface Gradio.
-    Ela executa o script torchrun em um subprocesso e transmite os logs em tempo real.
-    """
     if video_path is None:
         raise gr.Error("Por favor, faça o upload de um arquivo de vídeo ou imagem.")
-
-    # Cria diretórios temporários únicos para esta execução em /tmp
+    
     job_id = str(uuid.uuid4())
     input_dir = os.path.join("/tmp", "temp_inputs", job_id)
     output_dir = os.path.join("/tmp", "temp_outputs", job_id)
     os.makedirs(input_dir, exist_ok=True)
     os.makedirs(output_dir, exist_ok=True)
-
-    # O Gradio nos dá um caminho temporário. Copiamos o arquivo para nosso diretório de trabalho.
+    
     shutil.copy(video_path, input_dir)
-
+    
     log_output = ""
-
+    
+    # --- LÓGICA DO MONKEY PATCHING ---
+    patched_script_path = os.path.join("/tmp", f"inference_patched_{job_id}.py")
     try:
-        # O script de inferência é executado a partir de SEEDVR_DIR, então os caminhos precisam ser relativos
+        original_script_path = os.path.join(SEEDVR_DIR, "projects", "inference_seedvr2_3b.py")
+        with open(original_script_path, 'r') as f:
+            script_content = f.read()
+
+        # Define os caminhos hardcoded a serem substituídos
+        default_dit_path = "'./ckpts/seedvr2_ema_3b.pth'"
+        default_vae_path = "'./ckpts/ema_vae.pth'" # Assumindo que o VAE também é carregado assim no original
+
+        # Define os novos caminhos que apontam para nosso diretório em /tmp
+        patched_dit_path = f"'{os.path.join(CKPTS_DIR, 'seedvr2_ema_3b.pth')}'"
+        patched_vae_path = f"'{os.path.join(CKPTS_DIR, 'ema_vae.pth')}'"
+
+        # Aplica o "patch" substituindo o texto
+        script_content = script_content.replace(default_dit_path, patched_dit_path)
+        # Tenta substituir o caminho do VAE também, se existir
+        script_content = script_content.replace(default_vae_path, patched_vae_path)
+        
+        # Salva o script modificado em um arquivo temporário
+        with open(patched_script_path, 'w') as f:
+            f.write(script_content)
+        print(f"Script de inferência 'remendado' e salvo em: {patched_script_path}")
+        
+        # ------------------------------------
+
         input_folder_relative = os.path.relpath(input_dir, SEEDVR_DIR)
         output_folder_relative = os.path.relpath(output_dir, SEEDVR_DIR)
+        
+        # O torchrun deve executar o script remendado
+        patched_script_relative_path = os.path.relpath(patched_script_path, SEEDVR_DIR)
 
         command = [
             "torchrun", "--nproc-per-node=4",
-            "projects/inference_seedvr2_3b.py",
+            patched_script_relative_path, # <-- USA O SCRIPT MODIFICADO
             "--video_path", input_folder_relative,
             "--output_dir", output_folder_relative,
             "--seed", str(seed),
             "--res_h", str(res_h),
             "--res_w", str(res_w),
-            # Argumento crucial que diz ao script onde encontrar os modelos baixados
-            "--ckpt_dir", CKPTS_DIR
         ]
-
-        # Força o Python a não usar buffer de saída, garantindo logs em tempo real
+        
         env = os.environ.copy()
         env["PYTHONUNBUFFERED"] = "1"
-
-        log_output += f"Executando comando: {' '.join(command)}\n\n"
-        yield None, None, log_output # Limpa saídas antigas e mostra o comando na caixa de logs
-
+        
+        log_output += f"Executando comando com script remendado: {' '.join(command)}\n\n"
+        yield None, None, log_output
+        
         process = subprocess.Popen(
-            command,
-            cwd=SEEDVR_DIR,
-            stdout=subprocess.PIPE,
-            stderr=subprocess.STDOUT,
-            text=True,
-            encoding='utf-8',
-            env=env
+            command, cwd=SEEDVR_DIR, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
+            text=True, encoding='utf-8', env=env
         )
-
-        # Loop para capturar e transmitir a saída do subprocesso em tempo real para a UI
+        
         while True:
             output = process.stdout.readline()
-            if output == '' and process.poll() is not None:
-                break
+            if output == '' and process.poll() is not None: break
             if output:
                 log_output += output
-                # O yield atualiza a caixa de logs da interface Gradio
                 yield None, None, log_output
-
-        return_code = process.poll()
-        if return_code != 0:
-            raise gr.Error(f"A inferência falhou com o código de saída {return_code}. Verifique os logs para detalhes.")
+        
+        if process.poll() != 0:
+            raise gr.Error(f"A inferência falhou. Verifique os logs.")
 
         output_files = [f for f in os.listdir(output_dir) if f.endswith(('.mp4', '.png', '.jpg', '.jpeg'))]
         if not output_files:
-            raise gr.Error("A inferência foi concluída, mas nenhum arquivo de saída foi encontrado.")
-
+            raise gr.Error("Nenhum arquivo de saída foi encontrado.")
+            
         result_path = os.path.join(output_dir, output_files[0])
-
-        # Determina se a saída é imagem ou vídeo e a retorna para o componente correto
+        
         media_type, _ = mimetypes.guess_type(result_path)
         if media_type and media_type.startswith("image"):
             yield result_path, None, log_output
@@ -119,14 +122,14 @@ def run_inference(video_path, seed, res_h, res_w):
             yield None, result_path, log_output
 
     finally:
-        # Limpa o diretório de entrada temporário após a conclusão ou falha
         shutil.rmtree(input_dir, ignore_errors=True)
-        # O diretório de saída é mantido para que o Gradio possa exibir o resultado.
-        # Gradio gerencia a limpeza desses arquivos de saída.
+        # Limpa o script temporário que criamos
+        if os.path.exists(patched_script_path):
+            os.remove(patched_script_path)
 
-# --- Interface Gráfica Gradio ---
+# --- Interface Gráfica Gradio (sem alterações) ---
 with gr.Blocks(css="footer {display: none !important}") as demo:
-    gr.Markdown("# 🚀 Interface de Inferência para SeedVR2")
+    gr.Markdown("# 🚀 Interface de Inferência para SeedVR2 (com Monkey Patching)")
     gr.Markdown("Faça o upload de um vídeo ou imagem, ajuste os parâmetros e clique em 'Executar'.")
 
     with gr.Row():
@@ -137,18 +140,18 @@ with gr.Blocks(css="footer {display: none !important}") as demo:
                 res_h = gr.Number(value=720, label="Altura da Saída (res_h)")
                 res_w = gr.Number(value=1280, label="Largura da Saída (res_w)")
             run_button = gr.Button("Executar", variant="primary")
-
+        
         with gr.Column(scale=2):
             output_image = gr.Image(label="Saída de Imagem")
             output_video = gr.Video(label="Saída de Vídeo")
             log_box = gr.Textbox(label="Logs em Tempo Real", lines=15, autoscroll=True, interactive=False)
-
+            
     run_button.click(
         fn=run_inference,
         inputs=[input_media, seed, res_h, res_w],
         outputs=[output_image, output_video, log_box]
     )
-
+    
     
 
 demo.queue(max_size=10).launch()