Update app_seedvr.py

app_seedvr.py

@@ -1,213 +1,248 @@
-# app.py (Versão Corrigida)
+# app_seedvr.py
 
-import gradio as gr
-from PIL import Image
 import os
-import imageio
-from api.ltx_server import video_generation_service
-
-
-from huggingface_hub import logging
-
-
-logging.set_verbosity_error()
-logging.set_verbosity_warning()
-logging.set_verbosity_info()
-logging.set_verbosity_debug()
-
-enable_progress_bars()
-
-# --- FUNÇÕES DE AJUDA PARA A UI ---
-# ... (calculate_new_dimensions e handle_media_upload_for_dims permanecem as mesmas) ...
-TARGET_FIXED_SIDE = 768
-MIN_DIM_SLIDER = 256
-MAX_IMAGE_SIZE = 1280
-
-def calculate_new_dimensions(orig_w, orig_h):
-    if orig_w == 0 or orig_h == 0: return int(TARGET_FIXED_SIDE), int(TARGET_FIXED_SIDE)
-    if orig_w >= orig_h:
-        new_h, aspect_ratio = TARGET_FIXED_SIDE, orig_w / orig_h
-        new_w = round((new_h * aspect_ratio) / 32) * 32
-        new_w = max(MIN_DIM_SLIDER, min(new_w, MAX_IMAGE_SIZE))
-        new_h = max(MIN_DIM_SLIDER, min(new_h, MAX_IMAGE_SIZE)) 
-    else:
-        new_w, aspect_ratio = TARGET_FIXED_SIDE, orig_h / orig_w
-        new_h = round((new_w * aspect_ratio) / 32) * 32
-        new_h = max(MIN_DIM_SLIDER, min(new_h, MAX_IMAGE_SIZE))
-        new_w = max(MIN_DIM_SLIDER, min(new_w, MAX_IMAGE_SIZE))
-    return int(new_h), int(new_w)
-
-def handle_media_upload_for_dims(filepath, current_h, current_w):
-    if not filepath or not os.path.exists(str(filepath)): return gr.update(value=current_h), gr.update(value=current_w)
+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:
-        if str(filepath).lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
-            with Image.open(filepath) as img:
-                orig_w, orig_h = img.size
-        else: # Assumir que é um vídeo
-            with imageio.get_reader(filepath) as reader:
-                meta = reader.get_meta_data()
-                orig_w, orig_h = meta.get('size', (current_w, current_h))
-        new_h, new_w = calculate_new_dimensions(orig_w, orig_h)
-        return gr.update(value=new_h), gr.update(value=new_w)
+        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 processar mídia para dimensões: {e}")
-        return gr.update(value=current_h), gr.update(value=current_w)
-
-def update_frame_slider(duration):
-    """Atualiza o valor máximo do slider de frame do meio com base na duração."""
-    fps = 24.0
-    max_frames = int(duration * fps)
-    # Garante que o valor padrão não seja maior que o novo máximo
-    new_value = 48 if max_frames >= 48 else max_frames // 2
-    return gr.update(maximum=max_frames, value=new_value)
-
-
-# --- FUNÇÃO WRAPPER PARA CHAMAR O SERVIÇO ---
-def gradio_generate_wrapper(
-    prompt, negative_prompt, mode,
-    # Entradas de Keyframe
-    start_image,
-    middle_image, middle_frame, middle_weight,
-    end_image, end_weight,
-    # Outras entradas
-    input_video, height, width, duration,
-    frames_to_use, seed, randomize_seed,
-    guidance_scale, improve_texture,
+        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:
-        def progress_handler(step, total_steps):
-            progress(step / total_steps, desc="Salvando vídeo...")
-
-        output_path, used_seed = video_generation_service.generate(
-            prompt=prompt, negative_prompt=negative_prompt, mode=mode,
-            start_image_filepath=start_image,
-            middle_image_filepath=middle_image,
-            middle_frame_number=middle_frame,
-            middle_image_weight=middle_weight,
-            end_image_filepath=end_image,
-            end_image_weight=end_weight,
-            input_video_filepath=input_video,
-            height=int(height), width=int(width), duration=float(duration),
-            frames_to_use=int(frames_to_use), seed=int(seed),
-            randomize_seed=bool(randomize_seed), guidance_scale=float(guidance_scale),
-            improve_texture=bool(improve_texture), progress_callback=progress_handler
+        # 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 output_path, used_seed
-    except ValueError as e:
-        raise gr.Error(str(e))
+
     except Exception as e:
-        print(f"Erro inesperado na geração: {e}")
-        raise gr.Error("Ocorreu um erro inesperado. Verifique os logs.")
+        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)
+        )
 
-# --- DEFINIÇÃO DA INTERFACE GRADIO ---
-css = "#col-container { margin: 0 auto; max-width: 900px; }"
-with gr.Blocks(css=css) as demo:
-    gr.Markdown("# LTX Video com Keyframes")
-    gr.Markdown("Guie a geração de vídeo usando imagens de início, meio e fim.")
+# --- 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>
+        """
+    )
 
     with gr.Row():
-        with gr.Column():
-            with gr.Tab("image-to-video (Keyframes)") as image_tab:
-                i2v_prompt = gr.Textbox(label="Prompt", value="Uma bela transição entre as imagens", lines=2)
-                
-                with gr.Row():
-                    with gr.Column(scale=1):
-                        gr.Markdown("#### Início (Obrigatório)")
-                        start_image_i2v = gr.Image(label="Imagem de Início", type="filepath", sources=["upload", "clipboard"])
-                    with gr.Column(scale=1):
-                        gr.Markdown("#### Meio (Opcional)")
-                        middle_image_i2v = gr.Image(label="Imagem do Meio", type="filepath", sources=["upload", "clipboard"])
-                        middle_frame_i2v = gr.Slider(label="Frame Alvo", minimum=0, maximum=200, step=1, value=48)
-                        middle_weight_i2v = gr.Slider(label="Peso/Força", minimum=0.0, maximum=1.0, step=0.05, value=1.0)
-                    with gr.Column(scale=1):
-                        gr.Markdown("#### Fim (Opcional)")
-                        end_image_i2v = gr.Image(label="Imagem de Fim", type="filepath", sources=["upload", "clipboard"])
-                        end_weight_i2v = gr.Slider(label="Peso/Força", minimum=0.0, maximum=1.0, step=0.05, value=1.0)
+        # --- 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."
+                )
                 
-                i2v_button = gr.Button("Generate Image-to-Video", variant="primary")
-
-            with gr.Tab("text-to-video") as text_tab:
-                t2v_prompt = gr.Textbox(label="Prompt", value="A majestic dragon flying over a medieval castle", lines=3)
-                t2v_button = gr.Button("Generate Text-to-Video", variant="primary")
-
-            with gr.Tab("video-to-video") as video_tab:
-                video_v2v = gr.Video(label="Input Video", sources=["upload", "webcam"])
-                frames_to_use = gr.Slider(label="Frames to use from input video", minimum=9, maximum=257, value=9, step=8, info="Must be N*8+1.")
-                v2v_prompt = gr.Textbox(label="Prompt", value="Change the style to cinematic anime", lines=3)
-                v2v_button = gr.Button("Generate Video-to-Video", variant="primary")
-
-            duration_input = gr.Slider(label="Video Duration (seconds)", minimum=0.3, maximum=8.5, value=4, step=0.1)
-            improve_texture = gr.Checkbox(label="Improve Texture (multi-scale)", value=True, visible=True)
-
-        with gr.Column():
-            output_video = gr.Video(label="Generated Video", interactive=False)
-
-    with gr.Accordion("Advanced settings", open=False):
-        mode = gr.Dropdown(["text-to-video", "image-to-video", "video-to-video"], label="task", value="image-to-video", visible=False)
-        negative_prompt_input = gr.Textbox(label="Negative Prompt", value="worst quality, blurry, jittery", lines=2)
-        with gr.Row():
-            seed_input = gr.Number(label="Seed", value=42, precision=0)
-            randomize_seed_input = gr.Checkbox(label="Randomize Seed", value=True)
-        guidance_scale_input = gr.Slider(label="Guidance Scale (CFG)", minimum=1.0, maximum=10.0, value=3.0, step=0.1)
-        with gr.Row():
-            height_input = gr.Slider(label="Height", value=512, step=32, minimum=MIN_DIM_SLIDER, maximum=MAX_IMAGE_SIZE)
-            width_input = gr.Slider(label="Width", value=704, step=32, minimum=MIN_DIM_SLIDER, maximum=MAX_IMAGE_SIZE)
-
-    # --- LÓGICA DE EVENTOS DA UI ---
-    
-    start_image_i2v.upload(fn=handle_media_upload_for_dims, inputs=[start_image_i2v, height_input, width_input], outputs=[height_input, width_input])
-    video_v2v.upload(fn=handle_media_upload_for_dims, inputs=[video_v2v, height_input, width_input], outputs=[height_input, width_input])
-    duration_input.change(fn=update_frame_slider, inputs=duration_input, outputs=middle_frame_i2v)
+                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)*
+        """
+    )
     
-    image_tab.select(fn=lambda: "image-to-video", outputs=[mode])
-    text_tab.select(fn=lambda: "text-to-video", outputs=[mode])
-    video_tab.select(fn=lambda: "video-to-video", outputs=[mode])
+    # --- Lógica de Eventos da UI ---
     
-    # --- <INÍCIO DA CORREÇÃO> ---
-    # Reescrevendo as listas de inputs de forma explícita para evitar erros.
+    # 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]
+    )
     
-    # Placeholders para os botões que não usam certos inputs
-    none_image = gr.Textbox(visible=False, value=None)
-    none_video = gr.Textbox(visible=False, value=None)
-    
-    # Parâmetros comuns a todos
-    shared_params = [
-        height_input, width_input, duration_input, frames_to_use, 
-        seed_input, randomize_seed_input, guidance_scale_input, improve_texture
-    ]
-    
-    i2v_inputs = [
-        i2v_prompt, negative_prompt_input, mode,
-        start_image_i2v, middle_image_i2v, middle_frame_i2v, middle_weight_i2v,
-        end_image_i2v, end_weight_i2v,
-        none_video, # Placeholder para input_video
-        *shared_params
-    ]
-
-    t2v_inputs = [
-        t2v_prompt, negative_prompt_input, mode,
-        none_image, none_image, gr.Number(value=-1, visible=False), gr.Slider(value=0, visible=False), # Placeholders para keyframes
-        none_image, gr.Slider(value=0, visible=False),
-        none_video, # Placeholder para input_video
-        *shared_params
-    ]
-    
-    v2v_inputs = [
-        v2v_prompt, negative_prompt_input, mode,
-        none_image, none_image, gr.Number(value=-1, visible=False), gr.Slider(value=0, visible=False), # Placeholders para keyframes
-        none_image, gr.Slider(value=0, visible=False),
-        video_v2v, # Input de vídeo real
-        *shared_params
-    ]
-    
-    common_outputs = [output_video, seed_input]
-    
-    i2v_button.click(fn=gradio_generate_wrapper, inputs=i2v_inputs, outputs=common_outputs, api_name="image_to_video_keyframes")
-    t2v_button.click(fn=gradio_generate_wrapper, inputs=t2v_inputs, outputs=common_outputs, api_name="text_to_video")
-    v2v_button.click(fn=gradio_generate_wrapper, inputs=v2v_inputs, outputs=common_outputs, api_name="video_to_video")
-    # --- <FIM DA CORREÇÃO> ---
-
+    # 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.queue().launch(debug=True, share=False)
+    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
+    )