Update app.py

app.py

@@ -11,10 +11,64 @@
 # // 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 spaces
-import subprocess
 import os
 import sys
+import subprocess
+import importlib
+
+# --- ETAPA 0: CONFIGURAÇÃO ROBUSTA DO AMBIENTE COM CONDA ---
+
+CONDA_PYTHON_EXECUTABLE = "/home/user/miniconda3/bin/python"
+
+# Verifica se o ambiente conda já foi configurado
+if not os.path.exists(CONDA_PYTHON_EXECUTABLE):
+    print("Iniciando a configuração do ambiente Conda pela primeira vez. Isso pode levar alguns minutos...")
+    
+    # Função para executar comandos no shell
+    def run_command(command):
+        subprocess.run(command, shell=True, check=True, executable='/bin/bash')
+
+    # Instala o Miniconda
+    print("Baixando e instalando o Miniconda...")
+    run_command("wget https://repo.anaconda.com/miniconda/Miniconda3-py310_23.11.0-2-Linux-x86_64.sh -O miniconda.sh")
+    run_command("bash miniconda.sh -b -p /home/user/miniconda3")
+    run_command("rm miniconda.sh")
+    print("✅ Instalação do Miniconda concluída.")
+
+    # Define o caminho para o executável do Conda
+    conda_executable = "/home/user/miniconda3/bin/conda"
+    os.environ['PATH'] = f"/home/user/miniconda3/bin:{os.environ['PATH']}"
+
+    # Instala as dependências críticas com Conda (PyTorch, Apex pré-compilado)
+    print("Instalando PyTorch, Apex e dependências CUDA com Conda...")
+    run_command(f"{conda_executable} install -y -c pytorch -c nvidia pytorch torchvision pytorch-cuda=12.1 apex")
+    print("✅ Instalação de PyTorch e Apex via Conda concluída.")
+
+    # Instala as dependências restantes com o pip do ambiente conda
+    print("Instalando pacotes restantes com Pip (incluindo flash-attn)...")
+    pip_packages = [
+        "einops", "omegaconf", "lpips", "bson", "tensorflow", "opencv-python",
+        "sentencepiece", "diffusers", "rotary-embedding-torch", "transformers",
+        "tiktoken", "transformers_stream_generator", "torchmetrics", "pycocotools",
+        "torch-fidelity", "moviepy", "imageio", "tabulate", "deepdiff",
+        "parameterized", "mediapy", "av", "gradio",
+        "flash-attn" # Tenta instalar por último no ambiente já configurado
+    ]
+    run_command(f"{CONDA_PYTHON_EXECUTABLE} -m pip install {' '.join(pip_packages)} --no-build-isolation")
+    print("✅ Instalação das dependências do Pip concluída.")
+
+    print("\n*** Configuração do ambiente concluída. Reiniciando o script agora. ***\n")
+    # Executa novamente o script, mas desta vez com o interpretador Python do Conda
+    os.execle(CONDA_PYTHON_EXECUTABLE, CONDA_PYTHON_EXECUTABLE, *sys.argv)
+
+# A partir daqui, o script é executado pelo Python do ambiente Conda
+#---------------------------------------------------------------------
+
+import spaces
+from pathlib import Path
+from urllib.parse import urlparse
+from torch.hub import download_url_to_file
+import torch
 
 # --- ETAPA 1: Clonar o Repositório do GitHub ---
 repo_name = "SeedVR"
@@ -25,50 +79,10 @@ if not os.path.exists(repo_name):
 # --- ETAPA 2: Mudar para o Diretório e Configurar o Ambiente ---
 os.chdir(repo_name)
 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.")
-
-# --- ETAPA 3: Instalar Dependências Corretamente ---
-python_executable = sys.executable
-
-# Instalar flash-attn
-print("Instalando flash-attn...")
-subprocess.run(
-    [python_executable, "-m", "pip", "install", "flash-attn", "--no-build-isolation"],
-    check=True
-)
-print("✅ Configuração do flash-attn concluída.")
-
-# CORREÇÃO DEFINITIVA: Usar um método de instalação explícito para o Apex que força a compilação
-# das extensões CUDA através de variáveis de ambiente, que é a abordagem mais robusta.
-print("Clonando o repositório Apex para compilação manual...")
-if not os.path.exists("apex"):
-    subprocess.run("git clone 'https://huggingface.co/ByteDance-Seed/SeedVR2-3B/resolve/main/apex-0.1-cp310-cp310-linux_x86_64.whl'", shell=True, check=True)
-
-print("Instalando Apex a partir do código-fonte clonado com variáveis de ambiente de compilação...")
 
-# Cria uma cópia do ambiente atual e adiciona as flags de compilação do Apex
-env = os.environ.copy()
-env["APEX_CPP_EXT"] = "1"
-env["APEX_CUDA_EXT"] = "1"
-
-# Comando de instalação
-apex_install_command = [
-    python_executable, "-m", "pip", "install",
-    "--no-build-isolation",
-    "--no-cache-dir",
-    "./apex"
-]
-
-# Executa o comando com as variáveis de ambiente modificadas, o que força a compilação
-subprocess.run(apex_install_command, check=True, env=env)
-print("✅ Configuração do Apex concluída.")
-
-
-from pathlib import Path
-from urllib.parse import urlparse
-from torch.hub import download_url_to_file, get_dir
+# --- ETAPA 3: Baixar os Modelos Pré-treinados ---
+print("Baixando modelos pré-treinados...")
 
 def load_file_from_url(url, model_dir='.', progress=True, file_name=None):
     os.makedirs(model_dir, exist_ok=True)
@@ -81,10 +95,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
 
-# --- ETAPA 4: Baixar os Modelos Pré-treinados ---
-print("Baixando modelos pré-treinados...")
-import torch
-
 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',
@@ -98,12 +108,11 @@ for key, url in pretrain_model_url.items():
     load_file_from_url(url=url, model_dir=model_dir)
 
 
-# --- ETAPA 5: Executar a Aplicação Principal ---
+# --- ETAPA 4: Executar a Aplicação Principal ---
 import mediapy
 from einops import rearrange
 from omegaconf import OmegaConf
 import datetime
-from tqdm import tqdm
 import gc
 from PIL import Image
 import gradio as gr
@@ -127,10 +136,6 @@ os.environ["MASTER_PORT"] = "12355"
 os.environ["RANK"] = str(0)
 os.environ["WORLD_SIZE"] = str(1)
 
-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
-
 def configure_runner():
     config = load_config('configs_3b/main.yaml')
     runner = VideoDiffusionInfer(config)
@@ -161,7 +166,6 @@ 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()
-    # 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")]
@@ -169,7 +173,6 @@ def generation_loop(video_path, seed=666, fps_out=24):
     runner.configure_diffusion()
     set_seed(int(seed))
     os.makedirs("output", exist_ok=True)
-
     res_h, res_w = 1280, 720
     transform = Compose([
         NaResize(resolution=(res_h * res_w)**0.5, mode="area", downsample_only=False),
@@ -178,10 +181,8 @@ def generation_loop(video_path, seed=666, fps_out=24):
         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
@@ -189,14 +190,12 @@ 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 = [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
@@ -209,7 +208,7 @@ with gr.Blocks(title="SeedVR") as demo:
         <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 imagem e vídeo em um passo para conteúdo do mundo real e AIGC.
+        🔥 <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():
@@ -222,6 +221,5 @@ with gr.Blocks(title="SeedVR") as demo:
     output_video = gr.Video(label="Vídeo de Saída")
     download_link = gr.File(label="Baixar Resultado")
     run_button.click(fn=generation_loop, inputs=[input_file, seed, fps], outputs=[output_image, output_video, download_link])
-    
 
 demo.queue().launch(share=True)