# app_seedvr.py import os import sys from pathlib import Path from typing import Optional import gradio as gr import cv2 # --- 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 on_file_upload(file_obj): """ 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 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) ): """ 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. """ # 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) ) 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 log_buffer = ["▶ Starting inference process...\n"] last_log_message = "" was_input_video = _is_video(input_file_path) try: # 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) ) except Exception as e: error_message = f"❌ Inference failed: {e}" gr.Error(error_message) log_buffer.append(f"\n{error_message}") import traceback 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) ) # --- 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( """

📸 SeedVR - Image & Video Restoration 🚀

High-quality media upscaling powered by SeedVR-3B. Upload your file and see the magic.

""" ) 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." ) 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)* """ ) # --- Lógica de Eventos da UI --- # 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__": 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 )