Update app_seedvr.py

app_seedvr.py

@@ -1,273 +1,248 @@
-#!/usr/bin/env python3
-"""
-Standalone SeedVR2 Video Upscaler CLI Script
-(MODIFICADO PARA SER IMPORTÁVEL E SUPORTAR CALLBACKS DE PROGRESSO EM MULTIPROCESSING)
-"""
+# app_seedvr.py
 
-import sys
 import os
-import argparse
-import time
-import multiprocessing as mp
-import queue # Importa a classe de exceção para filas vazias
-
-# Garante o uso seguro de CUDA com multiprocessing, essencial para estabilidade.
-if mp.get_start_method(allow_none=True) != 'spawn':
-    mp.set_start_method('spawn', force=True)
-
-# -------------------------------------------------------------
-# 1) Configuração de alocação de memória da VRAM
-os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")
-
-# 2) Pré-análise dos argumentos para configurar a visibilidade dos dispositivos CUDA
-_pre_parser = argparse.ArgumentParser(add_help=False)
-_pre_parser.add_argument("--cuda_device", type=str, default=None)
-_pre_args, _ = _pre_parser.parse_known_args()
-if _pre_args.cuda_device is not None:
-    device_list_env = [x.strip() for x in _pre_args.cuda_device.split(',') if x.strip()!='']
-    if len(device_list_env) == 1:
-        os.environ["CUDA_VISIBLE_DEVICES"] = device_list_env[0]
-
-# -------------------------------------------------------------
-# 3) Importações pesadas (torch, etc.) são feitas após a configuração do ambiente.
-import torch
-import cv2
-import numpy as np
-from datetime import datetime
+import sys
 from pathlib import Path
-from src.utils.downloads import download_weight
+from typing import Optional
+import gradio as gr
+import cv2
 
-# Adiciona o diretório raiz do projeto ao path do sistema para permitir importações de `src`
-script_dir = os.path.dirname(os.path.abspath(__file__))
-if script_dir not in sys.path:
-    sys.path.insert(0, script_dir)
-root_dir = os.path.join(script_dir, '..', '..')
-if root_dir not in sys.path:
-    sys.path.insert(0, root_dir)
+# --- INTEGRAÇÃO COM A LÓGICA DO SERVIDOR ---
+try:
+    # Importa a classe SeedVRServer que agora atua como nossa biblioteca de inferência.
+    from api.seedvr_server import SeedVRServer
+except ImportError as e:
+    print(f"ERRO FATAL: Não foi possível importar o SeedVRServer. Detalhes: {e}")
+    # A aplicação não pode rodar sem a lógica do servidor.
+    raise
+
+# --- INICIALIZAÇÃO ---
+# Cria uma instância única e persistente do servidor.
+# A inicialização (clonar repo, baixar modelos) acontece apenas uma vez, no início.
+server = SeedVRServer()
+
+# --- FUNÇÕES AUXILIARES ---
+
+def _is_video(path: str) -> bool:
+    """Verifica se um caminho de arquivo corresponde a um tipo de vídeo."""
+    if not path: return False
+    import mimetypes
+    mime, _ = mimetypes.guess_type(path)
+    return (mime or "").startswith("video")
+
+def _extract_first_frame(video_path: str) -> Optional[str]:
+    """Extrai o primeiro frame de um vídeo e o salva como uma imagem JPG."""
+    if not video_path or not os.path.exists(video_path): return None
+    try:
+        vid_cap = cv2.VideoCapture(video_path)
+        if not vid_cap.isOpened():
+            print(f"Erro: Não foi possível abrir o vídeo em {video_path}")
+            return None
+        success, image = vid_cap.read()
+        vid_cap.release()
+        if not success:
+            print(f"Erro: Não foi possível ler o primeiro frame de {video_path}")
+            return None
+        
+        # Salva o frame no mesmo diretório do vídeo, com extensão .jpg
+        image_path = Path(video_path).with_suffix(".jpg")
+        cv2.imwrite(str(image_path), image)
+        return str(image_path)
+    except Exception as e:
+        print(f"Erro ao extrair o primeiro frame: {e}")
+        return None
 
-def extract_frames_from_video(video_path, debug=False, skip_first_frames=0, load_cap=None):
+def on_file_upload(file_obj):
     """
-    Extrai quadros de um vídeo e os converte para o formato de tensor.
+    Callback acionado quando o usuário faz o upload de um arquivo.
+    Verifica se o arquivo é um vídeo e sugere um `sp_size` apropriado.
     """
-    if debug: print(f"🎬 Extracting frames from video: {video_path}")
-    if not os.path.exists(video_path): raise FileNotFoundError(f"Video file not found: {video_path}")
-    cap = cv2.VideoCapture(video_path)
-    if not cap.isOpened(): raise ValueError(f"Cannot open video file: {video_path}")
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)); width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)); height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    if debug: print(f"📊 Video info: {frame_count} frames, {width}x{height}, {fps:.2f} FPS")
-    frames = []; frame_idx = 0; frames_loaded = 0
-    while True:
-        ret, frame = cap.read()
-        if not ret: break
-        if frame_idx < skip_first_frames: frame_idx += 1; continue
-        if load_cap is not None and load_cap > 0 and frames_loaded >= load_cap: break
-        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
-        frame = frame.astype(np.float32) / 255.0
-        frames.append(frame); frame_idx += 1; frames_loaded += 1
-    cap.release()
-    if len(frames) == 0: raise ValueError(f"No frames extracted from video: {video_path}")
-    if debug: print(f"✅ Extracted {len(frames)} frames")
-    frames_tensor = torch.from_numpy(np.stack(frames)).to(torch.float16)
-    if debug: print(f"📊 Frames tensor shape: {frames_tensor.shape}, dtype: {frames_tensor.dtype}")
-    return frames_tensor, fps
-
-def save_frames_to_video(frames_tensor, output_path, fps=30.0, debug=False):
+    if file_obj is None:
+        # Limpa os resultados e o log se o arquivo for removido
+        return gr.update(value=1), None, None, None, gr.update(value=None, visible=False)
+    
+    if _is_video(file_obj.name):
+        # Para vídeos, sugere um valor padrão para multi-GPU e torna o slider interativo
+        return gr.update(value=8, interactive=True), None, None, None, gr.update(value=None, visible=False)
+    else:
+        # Para imagens, trava o valor em 1
+        return gr.update(value=1, interactive=False), None, None, None, gr.update(value=None, visible=False)
+
+# --- FUNÇÃO PRINCIPAL DE INFERÊNCIA DA UI ---
+
+def run_inference_ui(
+    input_file_path: Optional[str],
+    resolution: str,
+    sp_size: int,
+    fps: float,
+    progress=gr.Progress(track_tqdm=True)
+):
     """
-    Salva um tensor de quadros em um arquivo de vídeo.
+    A função de callback principal do Gradio. Usa geradores (`yield`)
+    para permitir atualizações da UI em tempo real durante a tarefa de longa duração.
     """
-    if debug: print(f"🎬 Saving {frames_tensor.shape[0]} frames to video: {output_path}")
-    os.makedirs(os.path.dirname(output_path), exist_ok=True)
-    frames_np = (frames_tensor.cpu().numpy() * 255.0).astype(np.uint8)
-    T, H, W, C = frames_np.shape
-    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
-    out = cv2.VideoWriter(output_path, fourcc, fps, (W, H))
-    if not out.isOpened(): raise ValueError(f"Cannot create video writer for: {output_path}")
-    for i, frame in enumerate(frames_np):
-        frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
-        out.write(frame_bgr)
-    out.release()
-    if debug: print(f"✅ Video saved successfully: {output_path}")
+    # 1. Estado Inicial e Validação
+    # No início, desabilita o botão, limpa resultados anteriores e mostra a janela de log.
+    yield (
+        gr.update(interactive=False, value="Processing... 🚀"),
+        gr.update(value=None, visible=False),
+        gr.update(value=None, visible=False),
+        gr.update(value=None, visible=False),
+        gr.update(value="▶ Starting inference process...\n", visible=True)
+    )
 
-def _worker_process(proc_idx, device_id, frames_np, shared_args, return_queue, progress_queue=None):
-    """
-    Processo filho (worker) que executa o upscaling em uma GPU dedicada.
-    """
-    os.environ["CUDA_VISIBLE_DEVICES"] = str(device_id)
-    os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")
+    if not input_file_path:
+        gr.Warning("Please upload a media file first.")
+        # Reabilita o botão e esconde os componentes de saída
+        yield (gr.update(interactive=True, value="Restore Media"), None, None, None, gr.update(visible=False))
+        return
 
-    import torch
-    from src.core.model_manager import configure_runner
-    from src.core.generation import generation_loop
-    
-    frames_tensor = torch.from_numpy(frames_np).to(torch.float16)
-    
-    local_progress_callback = None
-    if progress_queue:
-        def callback_wrapper(batch_idx, total_batches, current_frames, message):
-            progress_queue.put((proc_idx, batch_idx, total_batches, message))
-        local_progress_callback = callback_wrapper
+    log_buffer = ["▶ Starting inference process...\n"]
+    last_log_message = ""
+    was_input_video = _is_video(input_file_path)
 
     try:
-        runner = configure_runner(shared_args["model"], shared_args["model_dir"], shared_args["preserve_vram"], shared_args["debug"])
-        result_tensor = generation_loop(
-            runner=runner, images=frames_tensor, cfg_scale=shared_args["cfg_scale"],
-            seed=shared_args["seed"], res_w=shared_args["res_w"], batch_size=shared_args["batch_size"],
-            preserve_vram=shared_args["preserve_vram"], temporal_overlap=shared_args["temporal_overlap"],
-            debug=shared_args["debug"],
-            progress_callback=local_progress_callback 
+        # Define um callback que será chamado pelo backend para atualizar o progresso e o log
+        def progress_callback_wrapper(step: float, desc: str):
+            """ Wrapper para formatar logs e atualizar o progresso. """
+            nonlocal last_log_message
+            # Só adiciona ao log se a mensagem for nova, para evitar poluição visual
+            if desc != last_log_message:
+                log_buffer.append(f"{desc}\n")
+                last_log_message = desc
+            # Atualiza o objeto de progresso do Gradio
+            progress(step, desc=desc)
+
+        # 2. Executa a Inferência
+        # Chama o método direto do servidor, passando o nosso callback.
+        video_result_path = server.run_inference_direct(
+            file_path=input_file_path,
+            seed=42, # Semente fixa conforme solicitado
+            res_h=int(resolution),
+            res_w=int(resolution), # Largura igual à altura
+            sp_size=int(sp_size),
+            fps=float(fps) if fps and fps > 0 else None,
+            progress=progress_callback_wrapper, # Passa nossa função de callback
+        )
+        
+        progress(1.0, desc="Complete!")
+        log_buffer.append("✅ Inference complete! Processing final output...\n")
+        
+        # 3. Processa e Exibe os Resultados
+        final_image, final_video = None, None
+        if was_input_video:
+            final_video = video_result_path
+            log_buffer.append("✅ Video result is ready.\n")
+        else: # Se a entrada foi uma imagem
+            final_image = _extract_first_frame(video_result_path)
+            final_video = video_result_path # Também disponibiliza o vídeo de 1 frame
+            log_buffer.append("✅ Image result extracted from video.\n")
+        
+        # Yield final para mostrar os resultados e reabilitar o botão
+        yield (
+            gr.update(interactive=True, value="Restore Media"),
+            gr.update(value=final_image, visible=final_image is not None),
+            gr.update(value=final_video, visible=final_video is not None),
+            gr.update(value=video_result_path, visible=video_result_path is not None),
+            ''.join(log_buffer)
         )
-        return_queue.put((proc_idx, result_tensor.cpu().numpy()))
+
     except Exception as e:
+        error_message = f"❌ Inference failed: {e}"
+        gr.Error(error_message)
+        log_buffer.append(f"\n{error_message}")
         import traceback
-        error_msg = f"ERROR in worker {proc_idx}: {e}\n{traceback.format_exc()}"
-        print(error_msg)
-        if progress_queue:
-            progress_queue.put((proc_idx, -1, -1, error_msg))
-        return_queue.put((proc_idx, error_msg))
-
-def _gpu_processing(frames_tensor, device_list, args, progress_callback=None):
-    """
-    Divide os quadros, gerencia os workers e monitora o progresso de forma robusta.
-    """
-    num_devices = len(device_list)
-    chunks = torch.chunk(frames_tensor, num_devices, dim=0)
-
-    manager = mp.Manager()
-    return_queue = manager.Queue()
-    progress_queue = manager.Queue() if progress_callback else None
-    workers = []
-
-    shared_args = {
-        "model": args.model, "model_dir": args.model_dir or "./models/SEEDVR2",
-        "preserve_vram": args.preserve_vram, "debug": args.debug, "cfg_scale": 1.0,
-        "seed": args.seed, "res_w": args.resolution, "batch_size": args.batch_size, "temporal_overlap": 0,
-    }
-
-    for idx, (device_id, chunk_tensor) in enumerate(zip(device_list, chunks)):
-        p = mp.Process(target=_worker_process, args=(idx, device_id, chunk_tensor.cpu().numpy(), shared_args, return_queue, progress_queue))
-        p.start()
-        workers.append(p)
-
-    results_np = [None] * num_devices
-    finished_workers = [False] * num_devices
-    worker_progress = [0.0] * num_devices
-    
-    while not all(finished_workers):
-        if progress_queue:
-            while not progress_queue.empty():
-                try:
-                    proc_idx, batch_idx, total_batches, message = progress_queue.get_nowait()
-                    if batch_idx == -1: # Mensagem de erro do worker
-                        raise RuntimeError(f"Worker {proc_idx} encontrou um erro: {message}")
-                    
-                    if total_batches > 0:
-                        worker_progress[proc_idx] = batch_idx / total_batches
-                    
-                    total_progress = sum(worker_progress) / num_devices
-                    progress_callback(total_progress, desc=f"GPU {proc_idx+1}/{num_devices}: {message}")
-                except queue.Empty:
-                    break
+        traceback.print_exc()
+        
+        # Yield para estado de erro: reabilita o botão e mostra o log com o erro
+        yield (
+            gr.update(interactive=True, value="Restore Media"),
+            None, None, None,
+            gr.update(value=''.join(log_buffer), visible=True)
+        )
 
-        while not return_queue.empty():
-            try:
-                proc_idx, result = return_queue.get_nowait()
-                if isinstance(result, str) and result.startswith("ERROR"):
-                     raise RuntimeError(f"Worker {proc_idx} falhou: {result}")
+# --- LAYOUT DA INTERFACE GRÁFICA (GRADIO) ---
+
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="SeedVR Media Restoration") as demo:
+    # Cabeçalho
+    gr.Markdown(
+        """
+        <div style='text-align: center; margin-bottom: 20px;'>
+            <h1>📸 SeedVR - Image & Video Restoration 🚀</h1>
+            <p>High-quality media upscaling powered by SeedVR-3B. Upload your file and see the magic.</p>
+        </div>
+        """
+    )
 
-                results_np[proc_idx] = result
-                worker_progress[proc_idx] = 1.0 # Marca como 100% concluído
-                finished_workers[proc_idx] = True
+    with gr.Row():
+        # --- Coluna da Esquerda: Entradas e Controles ---
+        with gr.Column(scale=1):
+            gr.Markdown("### 1. Upload Media")
+            input_media = gr.File(label="Input File (Video or Image)", type="filepath", interactive=True)
+            
+            gr.Markdown("### 2. Configure Settings")
+            with gr.Accordion("Generation Parameters", open=True):
+                resolution_select = gr.Dropdown(
+                    label="Resolution",
+                    choices=["480", "560", "720", "960", "1024", "2048"],
+                    value="480",
+                    info="Sets the output height and width to this value."
+                )
                 
-                total_progress = sum(worker_progress) / num_devices
-                if progress_callback:
-                    progress_callback(total_progress, desc=f"GPU {proc_idx+1}/{num_devices}: Concluído!")
-            except queue.Empty:
-                break
-
-        time.sleep(0.2)
-
-    for p in workers: p.join()
-
-    # Verifica se algum resultado está faltando, indicando um erro não capturado
-    if any(r is None for r in results_np):
-        raise RuntimeError("Um ou mais workers falharam em retornar um resultado.")
-
-    return torch.from_numpy(np.concatenate(results_np, axis=0)).to(torch.float16)
-
-def parse_arguments():
-    """Analisa os argumentos da linha de comando."""
-    parser = argparse.ArgumentParser(description="SeedVR2 Video Upscaler CLI")
-    parser.add_argument("--video_path", type=str, required=True, help="Path to input video file")
-    parser.add_argument("--seed", type=int, default=100, help="Random seed for generation (default: 100)")
-    parser.add_argument("--resolution", type=int, default=1072, help="Target resolution of the short side (default: 1072)")
-    parser.add_argument("--batch_size", type=int, default=5, help="Number of frames per batch (default: 5)")
-    parser.add_argument("--model", type=str, default="seedvr2_ema_3b_fp16.safetensors",
-                        choices=["seedvr2_ema_3b_fp16.safetensors", "seedvr2_ema_3b_fp8_e4m3fn.safetensors", 
-                                 "seedvr2_ema_7b_fp16.safetensors", "seedvr2_ema_7b_fp8_e4m3fn.safetensors"],
-                        help="Model to use")
-    parser.add_argument("--model_dir", type=str, default=None, help="Directory containing the model files")
-    parser.add_argument("--skip_first_frames", type=int, default=0, help="Skip the first frames during processing")
-    parser.add_argument("--load_cap", type=int, default=0, help="Maximum number of frames to load from video (default: load all)")
-    parser.add_argument("--output", type=str, default=None, help="Output path")
-    parser.add_argument("--output_format", type=str, default="video", choices=["video", "png"], help="Output format: 'video' (mp4) or 'png' images")
-    parser.add_argument("--preserve_vram", action="store_true", help="Enable VRAM preservation mode")
-    parser.add_argument("--debug", action="store_true", help="Enable debug logging")
-    parser.add_argument("--cuda_device", type=str, default=None, help="CUDA device id(s). e.g., '0' or '0,1' for multi-GPU")
-    
-    return parser.parse_args()
-
-def run_inference_logic(args, progress_callback=None):
-    """
-    Função principal que executa o pipeline de upscaling. Pode ser importada e chamada por outros scripts.
-    """
-    if args.debug:
-        print(f"📋 Argumentos da Lógica de Inferência: {vars(args)}")
-
-    if progress_callback: progress_callback(0.05, "Extracting frames...")
-    print("🎬 Extraindo frames do vídeo...")
-    start_time = time.time()
-    frames_tensor, original_fps = extract_frames_from_video(
-        args.video_path, args.debug, args.skip_first_frames, args.load_cap
+                sp_size_slider = gr.Slider(
+                    label="Frames per Batch (sp_size)",
+                    minimum=1, maximum=16, step=4, value=8,
+                    info="For multi-GPU videos. Automatically set to 1 for images."
+                )
+
+                fps_out = gr.Number(label="Output FPS (for Videos)", value=24, precision=0, info="Set to 0 to use the original FPS.")
+
+            run_button = gr.Button("Restore Media", variant="primary", icon="✨")
+
+        # --- Coluna da Direita: Resultados ---
+        with gr.Column(scale=2):
+            gr.Markdown("### 3. Results")
+            
+            # Janela de Log
+            log_window = gr.Textbox(
+                label="Inference Log 📝",
+                lines=8, max_lines=15,
+                interactive=False, visible=False, autoscroll=True
+            )
+
+            # Componentes de saída (começam invisíveis)
+            output_image = gr.Image(label="Image Result", show_download_button=True, type="filepath", visible=False)
+            output_video = gr.Video(label="Video Result", visible=False)
+            output_download = gr.File(label="Download Full Result (Video)", visible=False)
+
+    # --- Rodapé ---
+    gr.Markdown(
+        """
+        ---
+        *Space and Docker were developed by Carlex.*
+        *Contact: Email: Carlex22@gmail.com | GitHub: [carlex22](https://github.com/carlex22)*
+        """
     )
-    if args.debug:
-        print(f"🔄 Tempo de extração de frames: {time.time() - start_time:.2f}s")
-        
-    device_list = [d.strip() for d in str(args.cuda_device).split(',') if d.strip()] if args.cuda_device else ["0"]
-    if args.debug:
-        print(f"🚀 Usando dispositivos: {device_list}")
     
-    if progress_callback: progress_callback(0.1, "Starting generation...")
-    processing_start = time.time()
-    download_weight(args.model, args.model_dir)
+    # --- Lógica de Eventos da UI ---
     
-    result_tensor = _gpu_processing(frames_tensor, device_list, args, progress_callback)
-
-    generation_time = time.time() - processing_start
-    if args.debug:
-        print(f"🔄 Tempo de Geração: {generation_time:.2f}s")
-        print(f"📊 Resultado: {result_tensor.shape}, dtype: {result_tensor.dtype}")
-        
-    return result_tensor, original_fps, generation_time, len(frames_tensor)
-
-def main():
-    """
-    Função principal para execução via linha de comando (CLI).
-    """
-    print(f"🚀 SeedVR2 Video Upscaler CLI iniciado às {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
-    args = parse_arguments()
-    try:
-        result_tensor, original_fps, _, _ = run_inference_logic(args)
-        
-        print(f"💾 Salvando vídeo em: {args.output}")
-        save_frames_to_video(result_tensor, args.output, original_fps, args.debug)
-        print("✅ Upscaling via CLI concluído com sucesso!")
-        
-    except Exception as e:
-        print(f"❌ Erro durante o processamento via CLI: {e}")
-        import traceback
-        traceback.print_exc()
-        sys.exit(1)
+    # Ao fazer upload de um arquivo, ajusta o slider `sp_size` e limpa saídas antigas.
+    input_media.upload(
+        fn=on_file_upload,
+        inputs=[input_media],
+        outputs=[sp_size_slider, output_image, output_video, output_download, log_window]
+    )
+    
+    # Ao clicar no botão, executa a função de inferência principal.
+    run_button.click(
+        fn=run_inference_ui,
+        inputs=[input_media, resolution_select, sp_size_slider, fps_out],
+        outputs=[run_button, output_image, output_video, output_download, log_window],
+    )
 
 if __name__ == "__main__":
-    main()
+    demo.launch(
+        server_name=os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"),
+        server_port=int(os.getenv("GRADIO_SERVER_PORT", "7860")),
+        show_error=True
+    )