Update app_seedvr.py

app_seedvr.py

@@ -1,31 +1,38 @@
 # app_seedvr.py
 
 import os
+import sys
 from pathlib import Path
 from typing import Optional
 import gradio as gr
 import cv2
 
+# --- SERVER LOGIC INTEGRATION ---
+# This section ensures we can import and use the SeedVR engine directly.
 try:
-    # Importa a classe de servidor que agora é uma biblioteca local
+    # We need the SeedVRServer class which handles the inference logic.
     from api.seedvr_server import SeedVRServer
 except ImportError as e:
-    print(f"ERRO FATAL: Não foi possível importar o SeedVRServer. Detalhes: {e}")
-    # Se a importação falhar, a aplicação não pode continuar.
+    print(f"FATAL ERROR: Could not import SeedVRServer. Details: {e}")
+    # The application cannot run without the server logic.
     raise
 
-# Cria uma instância única do servidor. A inicialização (clonar repo, baixar modelos) acontece aqui.
+# --- INITIALIZATION ---
+# Create a single, persistent instance of the server.
+# This clones the repo and downloads models only once at startup.
 server = SeedVRServer()
 
+# --- HELPER FUNCTIONS ---
+
 def _is_video(path: str) -> bool:
-    """Verifica se um caminho de arquivo corresponde a um tipo de vídeo."""
+    """Checks if a file path corresponds to a video type."""
     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."""
+    """Extracts the first frame from a video and saves it as a JPG image."""
     if not video_path or not os.path.exists(video_path): return None
     try:
         vid_cap = cv2.VideoCapture(video_path)
@@ -34,106 +41,206 @@ def _extract_first_frame(video_path: str) -> Optional[str]:
         vid_cap.release()
         if not success: 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}")
+        print(f"Error extracting first frame: {e}")
         return None
 
-def ui_infer(
-    input_path: Optional[str],
-    seed: int, res_h: int, res_w: int,
-    sp_size: int, fps: float,
+def on_file_upload(file_obj):
+    """
+    Callback triggered when a user uploads a file.
+    It checks if the file is a video and suggests an appropriate `sp_size`.
+    """
+    if file_obj is None:
+        return 1 # Default to 1 if file is cleared
+    
+    if _is_video(file_obj.name):
+        # For videos, suggest a default value suitable for multi-GPU
+        return gr.update(value=4, interactive=True)
+    else:
+        # For images, lock the value to 1
+        return gr.update(value=1, interactive=False)
+
+# --- CORE INFERENCE FUNCTION ---
+
+def run_inference_ui(
+    input_file_path: Optional[str],
+    resolution: str,
+    sp_size: int,
+    fps: float,
     progress=gr.Progress(track_tqdm=True)
 ):
     """
-    Função de callback principal do Gradio. Agora chama a lógica de inferência diretamente.
+    The main callback function for Gradio. This is a generator (`yield`)
+    to allow for real-time UI updates during the long-running task.
     """
-    if not input_path:
-        gr.Warning("Por favor, faça o upload de um arquivo.")
-        return None, None, None
+    # 1. Initial State & Validation
+    # On start, disable the button, clear previous results, and make the log visible.
+    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="Waiting for logs...", visible=True)
+    )
+
+    if not input_file_path:
+        gr.Warning("Please upload a media file first.")
+        # Re-enable button and hide outputs
+        yield (
+            gr.update(interactive=True, value="Restore Media"),
+            None, None, None, gr.update(visible=False)
+        )
+        return
+
+    # Use a simple list to act as a log buffer that can be updated by a callback
+    log_buffer = ["▶ Starting inference process...\n"]
+    yield gr.update(), None, None, None, ''.join(log_buffer)
+
+    def progress_callback(step: float, desc: str):
+        """A simple callback to append messages to our log buffer."""
+        # This function can be passed to the backend if it supports it.
+        # For now, we'll call it manually from this UI function.
+        log_buffer.append(f"⏳ [{int(step*100)}%] {desc}\n")
+        progress.update(amount=step, desc=desc)
+
+    was_input_video = _is_video(input_file_path)
 
-    was_input_video = _is_video(input_path)
-    
     try:
-        # Desabilita o botão enquanto processa
-        yield gr.update(interactive=False, value="Processando..."), None, None, None
+        # 2. Execute Inference
+        progress_callback(0.1, "Calling backend engine...")
+        yield gr.update(), None, None, None, ''.join(log_buffer)
 
-        # Chama o método direto do servidor, passando o objeto de progresso do Gradio
+        # Call the server's direct inference method. This is a blocking call.
         video_result_path = server.run_inference_direct(
-            file_path=input_path,
-            seed=int(seed),
-            res_h=int(res_h),
-            res_w=int(res_w),
+            file_path=input_file_path,
+            seed=42, # Using a fixed seed as requested
+            res_h=int(resolution),
+            res_w=int(resolution), # Set width equal to height
             sp_size=int(sp_size),
             fps=float(fps) if fps and fps > 0 else None,
-            progress=progress,
+            progress=progress, # Pass the Gradio progress object
         )
         
-        progress(1.0, desc="Concluído!")
-
+        progress_callback(1.0, "Inference complete! Processing final output...")
+        yield gr.update(), None, None, None, ''.join(log_buffer)
+        
+        # 3. Process and Display Results
         final_image, final_video = None, None
         if was_input_video:
             final_video = video_result_path
-        else: # Se a entrada foi uma imagem
+            log_buffer.append(f"✅ Video result is ready.\n")
+        else: # If input was an image
             final_image = _extract_first_frame(video_result_path)
-            final_video = video_result_path
-            
-        # Retorna o resultado e reabilita o botão
+            final_video = video_result_path # Also provide the 1-frame video
+            log_buffer.append(f"✅ Image result extracted from video.\n")
+        
+        # Final yield to show the results and re-enable the button
         yield (
-            gr.update(interactive=True, value="Restaurar Mídia"),
+            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)
+            gr.update(value=video_result_path, visible=video_result_path is not None),
+            ''.join(log_buffer)
         )
 
     except Exception as e:
-        error_message = f"A inferência falhou: {e}"
+        error_message = f"❌ Inference failed: {e}"
         gr.Error(error_message)
         print(error_message)
         import traceback
         traceback.print_exc()
-        # Limpa os resultados e reabilita o botão em caso de erro
-        yield gr.update(interactive=True, value="Restaurar Mídia"), None, None, None
-
-# --- Construção da Interface Gráfica ---
-with gr.Blocks(title="SeedVR (Aduc-SDR)", theme=gr.themes.Soft()) as demo:
-    gr.HTML("""
-        <div style='text-align:center; margin-bottom: 20px;'>
-            <h1>SeedVR - Restauração de Imagem e Vídeo</h1>
-            <p>Implementação com backend Aduc-SDR</p>
+        
+        # Yield an error state and re-enable the button
+        yield (
+            gr.update(interactive=True, value="Restore Media"),
+            None, None, None,
+            gr.update(value=f"{''.join(log_buffer)}\n{error_message}", visible=True)
+        )
+
+
+# --- GRADIO UI LAYOUT ---
+
+with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="SeedVR Media Restoration") as demo:
+    # Header
+    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():
+        # --- Left Column: Inputs & Controls ---
         with gr.Column(scale=1):
-            inp = gr.File(label="Arquivo de Entrada (Vídeo .mp4 ou Imagem)", type="filepath")
+            gr.Markdown("### 1. Upload Media")
+            input_media = gr.File(label="Input File (Video or Image)", type="filepath")
             
-            with gr.Accordion("Parâmetros de Geração", open=True):
-                with gr.Row():
-                    seed = gr.Number(label="Seed", value=42, precision=0)
-                    fps_out = gr.Number(label="FPS de Saída (para Vídeos)", value=24, precision=0, info="0 para usar o FPS original.")
-                with gr.Row():
-                    res_h = gr.Number(label="Altura (Height)", value=720, precision=0)
-                    res_w = gr.Number(label="Largura (Width)", value=1280, precision=0)
-
-                sp_size = gr.Slider(label="Paralelismo de Sequência (sp_size)", minimum=1, maximum=160, step=4, value=4, info="Para vídeos em multi-GPU. Use 1 para imagens.")
-
-            run_button = gr.Button("Restaurar Mídia", variant="primary")
-
+            gr.Markdown("### 2. Configure Settings")
+            with gr.Accordion("Generation Parameters", open=True):
+                resolution_select = gr.Dropdown(
+                    label="Resolution (Short Edge)",
+                    choices=["480", "560", "720", "960", "1024"],
+                    value="480",
+                    info="The output height and width will be set to this value."
+                )
+                
+                sp_size_slider = gr.Slider(
+                    label="Sequence Parallelism (sp_size)",
+                    minimum=1, maximum=16, step=1, value=4,
+                    info="For multi-GPU videos. This will be 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="✨")
+
+        # --- Right Column: Outputs ---
         with gr.Column(scale=2):
-            gr.Markdown("### Resultado")
-            out_image = gr.Image(label="Resultado (Imagem)", show_download_button=True, type="filepath", visible=True)
-            out_video = gr.Video(label="Resultado (Vídeo)")
-            out_download = gr.File(label="Baixar Resultado (Vídeo)")
-
-    # A função click agora é um gerador.
+            gr.Markdown("### 3. Results")
+            
+            # Log window
+            log_window = gr.Textbox(
+                label="Inference Log 📝",
+                lines=8,
+                max_lines=15,
+                interactive=False,
+                visible=False, # Starts hidden
+                autoscroll=True,
+            )
+
+            # Output components start hidden and are made visible upon completion
+            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)
+
+    # --- Footer ---
+    gr.Markdown(
+        """
+        ---
+        *Space and Docker were developed by Carlex.*
+        *Contact: Email: Carlex22@gmail.com | GitHub: [carlex22](https://github.com/carlex22)*
+        """
+    )
+    
+    # --- Event Handlers ---
+    
+    # When a file is uploaded, automatically adjust the sp_size slider
+    input_media.upload(
+        fn=on_file_upload,
+        inputs=[input_media],
+        outputs=[sp_size_slider]
+    )
+    
+    # When the "Restore Media" button is clicked, run the main inference function
     run_button.click(
-        fn=ui_infer,
-        inputs=[inp, seed, res_h, res_w, sp_size, fps_out],
-        outputs=[run_button, out_image, out_video, out_download],
+        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__":