Update app_seedvr.py

app_seedvr.py

@@ -1,81 +1,290 @@
-#!/usr/bin/env python3
-# -*- coding: utf-8 -*-
-
 import os
-from pathlib import Path
-from typing import Optional, Tuple
+import sys
+import uuid
+import json
+from typing import List, Dict, Any, Optional
+
+import torch
 import gradio as gr
+from PIL import Image
+from omegaconf import OmegaConf, DictConfig
 
-from services.seedvr import SeedVRRefineService
+# --- 1. CONFIGURAÇÃO E IMPORTS ---
 
-APP_HOME = Path(os.environ.get("APP_HOME", "/app"))
-OUT_DIR = APP_HOME / "outputs" / "seedvr_refine"
+# Adiciona o diretório do VINCIE ao Python Path para importação de módulos.
+VINCIE_DIR = os.getenv("VINCIE_DIR", "/app/VINCIE")
+if VINCIE_DIR not in sys.path:
+    sys.path.append(VINCIE_DIR)
 
-svc = SeedVRRefineService()
+try:
+    from generate import VINCIEGenerator
+    from common.config import load_config
+    from common.seed import shift_seed
+except ImportError:
+    print(f"FATAL: Não foi possível importar os módulos do VINCIE. "
+          f"Verifique se o repositório está em '{VINCIE_DIR}'.")
+    raise
 
-def setup() -> str:
-    try:
-        token = os.environ.get("HF_TOKEN") or os.environ.get("HFTOKEN")
-        msg = svc.ensure_model(max_workers=int(os.environ.get("MAX_WORKERS", "48")), token=token)
-        svc.ensure_apex(enable_shim=True)
-        return f"Setup concluído: {msg} (apex shim ativo)."
-    except Exception as e:
-        return f"Setup falhou: {e}"
-
-def refine_ui(
-    input_file: gr.File,
-    upscale: float,
-    strength: float,
-    denoise: float,
-    t_consistency: float,
-    fps_out: int,
-    tile: int,
-    dtype: str
-) -> Tuple[str, str]:
+# --- 2. INICIALIZAÇÃO DO MODELO (SINGLETON) ---
+
+MODEL: Optional[VINCIEGenerator] = None
+DEVICE: Optional[torch.device] = None
+
+def setup_model():
+    """
+    Inicializa e configura o modelo VINCIE em uma única GPU.
+    Esta função é chamada uma vez no início da aplicação.
+    """
+    global MODEL, DEVICE
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("FATAL: Nenhuma GPU compatível com CUDA foi encontrada.")
+
+    num_gpus = torch.cuda.device_count()
+    if num_gpus == 0:
+        raise RuntimeError("FATAL: Nenhuma GPU foi detectada pelo PyTorch.")
+    
+    print(f"INFO: Detectadas {num_gpus} GPUs. A aplicação usará 'cuda:0'.")
+    DEVICE = torch.device("cuda:0")
+    torch.cuda.set_device(DEVICE)
+
+    config_path = os.path.join(VINCIE_DIR, "configs/generate.yaml")
+    print(f"INFO: Carregando e resolvendo configuração de '{config_path}'...")
+    config = load_config(config_path, [])
+    
+    print("INFO: Instanciando VINCIEGenerator...")
+    model_instance = VINCIEGenerator(config)
+    
+    print("INFO: Executando sequência de inicialização do VINCIE...")
+    model_instance.configure_persistence()
+    model_instance.configure_models()
+    model_instance.configure_diffusion()
+    
+    if not hasattr(model_instance, 'dit'):
+        raise RuntimeError("FATAL: Falha ao inicializar o componente DiT do modelo.")
+
+    # Move todos os componentes para o dispositivo principal.
+    model_instance.dit.to(DEVICE)
+    model_instance.vae.to(DEVICE)
+    model_instance.text_encoder.to(DEVICE)
+    
+    MODEL = model_instance
+    print(f"✅ SUCESSO: Modelo VINCIE pronto para uso na GPU {DEVICE}.")
+
+# --- 3. LÓGICAS DE INFERÊNCIA ---
+
+def _execute_vincie_logic(
+    prompt_config: DictConfig,
+    steps: int,
+    cfg_scale: float,
+    seed: int,
+    pad_img_placeholder: bool,
+    resolution: int
+) -> Image.Image:
+    """
+    Função central que executa a pipeline de inferência do VINCIE.
+    """
+    # Salva o estado original da configuração para restaurá-lo depois.
+    original_config_state = {
+        "steps": MODEL.config.diffusion.timesteps.sampling.steps,
+        "seed": MODEL.config.generation.seed,
+        "pad": MODEL.config.generation.pad_img_placehoder,
+        "resolution": MODEL.config.generation.resolution,
+    }
+    
     try:
-        inp = Path(input_file.name) if hasattr(input_file, "name") else Path(input_file)
-        out = svc.refine(
-            input_path=inp,
-            output_dir=OUT_DIR,
-            upscale=upscale,
-            strength=strength,
-            denoise=denoise,
-            t_consistency=t_consistency,
-            fps_out=(None if fps_out <= 0 else fps_out),
-            tile=(None if tile <= 0 else tile),
-            dtype=(None if not dtype else dtype)
+        OmegaConf.set_readonly(MODEL.config, False)
+        
+        # 1. Aplica configurações dinâmicas da UI.
+        MODEL.config.diffusion.timesteps.sampling.steps = int(steps)
+        MODEL.configure_diffusion()  # Recria o sampler com os novos passos.
+        
+        current_seed = seed if seed != -1 else torch.randint(0, 2**32 - 1, (1,)).item()
+        MODEL.config.generation.seed = shift_seed(current_seed, 0)
+        MODEL.config.generation.pad_img_placehoder = pad_img_placeholder
+        MODEL.config.generation.resolution = int(resolution)
+
+        # Log detalhado dos argumentos que serão enviados para a pipeline.
+        _log_pipeline_args(prompt_config, steps, cfg_scale, MODEL.config.generation.seed, resolution, pad_img_placeholder)
+
+        # 2. Prepara as entradas para o modelo.
+        text_pos, condition, noise, _, _ = MODEL.prepare_input(
+            prompt=prompt_config, repeat_idx=0, device=DEVICE
         )
-        return (f"OK: saída em {out}", out)
-    except Exception as e:
-        return (f"Erro: {e}", "")
-
-with gr.Blocks(title="SeedVR — Refinamento (sem Apex)", theme="Default") as demo:
-    gr.Markdown("### SeedVR — Refinamento de vídeo/imagem (qualidade, sem Apex)")
-    with gr.Row():
-        with gr.Column(scale=2):
-            input_file = gr.File(label="Vídeo/Imagem de entrada", type="filepath")
-            upscale = gr.Slider(1.0, 2.0, value=1.0, step=0.1, label="Upscale factor")
-            strength = gr.Slider(0.0, 1.0, value=0.35, step=0.05, label="Força do refine")
-            denoise = gr.Slider(0.0, 1.0, value=0.1, step=0.05, label="Denoise")
-            tcons = gr.Slider(0.0, 1.0, value=0.7, step=0.05, label="Consistência temporal")
-            fps_out = gr.Number(value=0, precision=0, label="FPS saída (0 = manter)")
-            tile = gr.Number(value=0, precision=0, label="Tile (0 = off)")
-            dtype = gr.Dropdown(choices=["", "bfloat16", "float16"], value="", label="Dtype override")
-            go = gr.Button("Refinar", variant="primary")
-        with gr.Column(scale=1):
-            status = gr.Textbox(label="Status")
-            outdir = gr.Textbox(label="Saída", value=str(OUT_DIR), interactive=False)
-
-    go.click(
-        fn=refine_ui,
-        inputs=[input_file, upscale, strength, denoise, tcons, fps_out, tile, dtype],
-        outputs=[status, outdir]
-    )
-
-    init = gr.Button("Preparar modelos (setup)")
-    init.click(fn=setup, outputs=status)
+        
+        # 3. Executa a inferência.
+        with torch.no_grad():
+            samples = MODEL.inference(
+                noises=[noise],
+                conditions=[condition],
+                texts_pos=[text_pos],
+                texts_neg=[MODEL.config.generation.negative_prompt],
+                cfg_scale=cfg_scale
+            )
+            
+        if not samples:
+            raise RuntimeError("A inferência do modelo não produziu resultados.")
+            
+        # 4. Processa a saída para formato de imagem.
+        output_tensor = samples[0][:, -1, :, :]
+        output_image_np = output_tensor.clip(-1, 1).add(1).div(2).mul(255).byte().permute(1, 2, 0).cpu().numpy()
+        return Image.fromarray(output_image_np)
+
+    finally:
+        # 5. Restaura a configuração original para garantir consistência entre chamadas.
+        OmegaConf.set_readonly(MODEL.config, False)
+        for key, value in original_config_state.items():
+            if key == "steps":
+                MODEL.config.diffusion.timesteps.sampling.steps = value
+            else:
+                OmegaConf.update(MODEL.config.generation, key, value, merge=True)
+        OmegaConf.set_readonly(MODEL.config, True)
+        MODEL.configure_diffusion() # Restaura o sampler padrão.
+
+def run_single_turn_inference(
+    input_image: str, prompt: str, aspect_ratio: str, resolution: int, steps: int, cfg_scale: float, seed: int
+) -> Image.Image:
+    """Handler para a aba 'Edição Simples'."""
+    if not all([input_image, prompt]):
+        raise gr.Error("É necessário fornecer uma imagem de entrada e um prompt.")
+    
+    prompt_config = OmegaConf.create({
+        "index": 0, "img_paths": [input_image], "context": [prompt], "aspect_ratio": aspect_ratio
+    })
+    
+    return _execute_vincie_logic(prompt_config, steps, cfg_scale, seed, pad_img_placeholder=True, resolution=resolution)
+
+def run_multi_turn_inference(
+    input_image: str, prompts_text: str, steps: int, cfg_scale: float, seed: int, progress=gr.Progress()
+) -> List[Image.Image]:
+    """Handler para a aba 'Edição em Múltiplos Turnos'."""
+    if not all([input_image, prompts_text]):
+        raise gr.Error("É necessário fornecer uma imagem de entrada e pelo menos um prompt.")
+    
+    prompts = [p.strip() for p in prompts_text.splitlines() if p.strip()]
+    if not prompts:
+        raise gr.Error("Nenhum prompt válido fornecido.")
+        
+    output_images_with_paths = []
+    
+    for i, prompt in enumerate(progress.tqdm(prompts, desc="Processando turnos")):
+        image_paths = [input_image] + [path for _, path in output_images_with_paths]
+        context_prompts = prompts[:i+1]
+        
+        prompt_config = OmegaConf.create({
+            "index": i, "img_paths": image_paths, "context": context_prompts, "aspect_ratio": "keep_ratio"
+        })
+        
+        turn_seed = seed if seed == -1 else seed + i
+        result_image = _execute_vincie_logic(prompt_config, steps, cfg_scale, turn_seed, pad_img_placeholder=True, resolution=512)
+        
+        temp_path = os.path.join("/tmp", f"{uuid.uuid4()}.png")
+        result_image.save(temp_path)
+        output_images_with_paths.append((result_image, temp_path))
+
+    return [img for img, _ in output_images_with_paths]
+
+def run_multi_concept_inference(prompt: str, *images: str) -> Image.Image:
+    """Handler para a aba 'Composição de Conceitos'."""
+    image_paths = [img for img in images if img is not None]
+    if not image_paths or not prompt.strip():
+        raise gr.Error("É necessário um prompt e pelo menos uma imagem de entrada.")
+
+    # Constrói a lista de prompts: N-1 placeholders + 1 prompt de composição.
+    prompts_list = [f"<IMG{i}>: " for i in range(1, len(image_paths))]
+    prompts_list.append(prompt)
+        
+    prompt_config = OmegaConf.create({
+        "index": 0, "img_paths": image_paths, "context": prompts_list, "aspect_ratio": "1:1"
+    })
+    
+    # Usa parâmetros fixos para esta funcionalidade, conforme documentação do VINCIE.
+    return _execute_vincie_logic(prompt_config, steps=50, cfg_scale=7.5, seed=1, pad_img_placeholder=False, resolution=512)
+
+def _log_pipeline_args(prompt_config, steps, cfg_scale, final_seed, resolution, pad_placeholder):
+    """Função auxiliar para imprimir os argumentos exatos enviados à pipeline do VINCIE."""
+    log_data = {
+        "--- INÍCIO DOS ARGUMENTOS DA PIPELINE VINCIE ---": "",
+        "1. Configuração do Prompt": OmegaConf.to_container(prompt_config, resolve=True),
+        "2. Parâmetros de Difusão": {
+            "steps": int(steps),
+            "cfg_scale": float(cfg_scale),
+        },
+        "3. Parâmetros de Geração": {
+            "seed_final": int(final_seed),
+            "resolution": int(resolution),
+            "pad_img_placeholder": bool(pad_placeholder),
+        },
+        "--- FIM DOS ARGUMENTOS ---": ""
+    }
+    print(json.dumps(log_data, indent=2, ensure_ascii=False))
+
+# --- 4. CONSTRUÇÃO DA INTERFACE GRADIO ---
+
+def create_ui():
+    """Cria e retorna a interface Gradio completa com todas as abas e controles."""
+    with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="VINCIE Playground") as demo:
+        gr.Markdown("# 🖼️ **VINCIE Playground**\nExplore as diferentes capacidades do modelo VINCIE.")
+
+        # Controles avançados compartilhados pelas abas 1 e 2
+        with gr.Accordion("Opções Avançadas (para Abas 1 e 2)", open=False):
+            steps_input = gr.Slider(label="Passos de Inferência", minimum=10, maximum=100, step=1, value=50)
+            cfg_scale_input = gr.Slider(label="Escala de Orientação (CFG)", minimum=1.0, maximum=15.0, step=0.5, value=7.5)
+            seed_input = gr.Number(label="Semente (Seed)", value=-1, precision=0, info="Use -1 para aleatório.")
+
+        with gr.Tabs():
+            # Aba 1: Edição Simples
+            with gr.TabItem("Edição Simples"):
+                with gr.Row(equal_height=False):
+                    with gr.Column(scale=1):
+                        single_turn_img_in = gr.Image(type="filepath", label="Imagem de Entrada")
+                        single_turn_prompt = gr.Textbox(lines=2, label="Prompt de Edição")
+                        with gr.Accordion("Opções de Imagem", open=True):
+                             aspect_ratio_input = gr.Dropdown(label="Aspect Ratio", choices=["keep_ratio", "1:1", "16:9", "9:16", "4:3", "3:4"], value="keep_ratio")
+                             resolution_input = gr.Slider(label="Resolução (lado menor)", minimum=256, maximum=1024, step=64, value=512)
+                        single_turn_button = gr.Button("Gerar", variant="primary")
+                    with gr.Column(scale=1):
+                        single_turn_img_out = gr.Image(label="Resultado", interactive=False)
+                gr.Examples([["/app/VINCIE/assets/woman_pineapple.png", "Adicione uma coroa na cabeça da mulher."]], [single_turn_img_in, single_turn_prompt])
+
+            # Aba 2: Edição em Múltiplos Turnos
+            with gr.TabItem("Edição em Múltiplos Turnos"):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        multi_turn_img_in = gr.Image(type="filepath", label="Imagem de Entrada Inicial")
+                        multi_turn_prompts = gr.Textbox(lines=5, label="Prompts (um por linha)", placeholder="Turno 1: faça isso\nTurno 2: agora mude aquilo...")
+                        multi_turn_button = gr.Button("Gerar Sequência", variant="primary")
+                    with gr.Column(scale=2):
+                        multi_turn_gallery_out = gr.Gallery(label="Resultados dos Turnos", columns=3, height="auto")
+
+            # Aba 3: Composição de Conceitos
+            with gr.TabItem("Composição de Conceitos"):
+                gr.Markdown("Faça o upload de até 6 imagens (`<IMG0>` a `<IMG5>`) e escreva um prompt que as combine para gerar uma nova imagem (`<IMG6>`).")
+                with gr.Row():
+                    concept_inputs = [gr.Image(type="filepath", label=f"Imagem {i} (<IMG{i}>)") for i in range(6)]
+                concept_prompt = gr.Textbox(lines=4, label="Prompt de Composição Final", value="Baseado em <IMG0> e <IMG1>, um retrato do homem de <IMG0> usando o chapéu de <IMG1>. Saída <IMG6>:")
+                concept_button = gr.Button("Compor Imagem", variant="primary")
+                concept_img_out = gr.Image(label="Resultado da Composição", interactive=False)
+
+        # Conecta os botões às suas respectivas funções de backend
+        single_turn_button.click(fn=run_single_turn_inference, inputs=[single_turn_img_in, single_turn_prompt, aspect_ratio_input, resolution_input, steps_input, cfg_scale_input, seed_input], outputs=[single_turn_img_out])
+        multi_turn_button.click(fn=run_multi_turn_inference, inputs=[multi_turn_img_in, multi_turn_prompts, steps_input, cfg_scale_input, seed_input], outputs=[multi_turn_gallery_out])
+        concept_button.click(fn=run_multi_concept_inference, inputs=[concept_prompt] + concept_inputs, outputs=[concept_img_out])
+        
+    return demo
+
+# --- 5. PONTO DE ENTRADA DA APLICAÇÃO ---
 
 if __name__ == "__main__":
-    port = int(os.environ.get("PORT", "7860"))
-    demo.launch(server_name="0.0.0.0", server_port=port, allowed_paths=[str(OUT_DIR), str(APP_HOME / "ckpt" / "SeedVR2-3B")])
-    
+    setup_model()
+    ui = create_ui()
+    
+    server_name = os.environ.get("GRADIO_SERVER_NAME", "0.0.0.0")
+    server_port = int(os.environ.get("GRADIO_SERVER_PORT", 7860))
+    enable_queue = os.environ.get("GRADIO_ENABLE_QUEUE", "True").lower() == "true"
+
+    print(f"INFO: Lançando a interface Gradio em http://{server_name}:{server_port}")
+    if enable_queue:
+        ui.queue()
+
+    ui.launch(server_name=server_name, server_port=server_port)