Update app_vince.py

app_vince.py

@@ -1,33 +1,39 @@
 import os
 import sys
+import uuid
+from typing import List, Dict, Any, Optional, Tuple
+
 import torch
 import gradio as gr
 from PIL import Image
-from omegaconf import OmegaConf
+from omegaconf import OmegaConf, DictConfig
 
-# --- ETAPA 1: CONFIGURAÇÃO DO AMBIENTE ---
+# --- 1. CONFIGURAÇÃO E IMPORTS ---
 
+# Adiciona o diretório do VINCIE ao path do Python para permitir a importação de seus módulos.
 VINCIE_DIR = "/app/VINCIE"
 if VINCIE_DIR not in sys.path:
     sys.path.append(VINCIE_DIR)
 
 try:
     from generate import VINCIEGenerator
-    from common.distributed import get_device
-    from common.config import load_config 
+    from common.config import load_config
+    from common.seed import shift_seed
 except ImportError as e:
-    print(f"FATAL: Não foi possível importar os módulos do VINCIE.")
-    print(f"Verifique se o diretório '{VINCIE_DIR}' contém o código-fonte clonado.")
+    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 e
 
-# --- ETAPA 2: INICIALIZAÇÃO DO MODELO (Simplificado para uma GPU) ---
+# --- 2. INICIALIZAÇÃO DO MODELO (SINGLETON) ---
 
-MODEL: VINCIEGenerator = None
-DEVICE: torch.device = None
+# Variáveis globais para armazenar a instância única do modelo e o dispositivo.
+MODEL: Optional[VINCIEGenerator] = None
+DEVICE: Optional[torch.device] = None
 
 def setup_model():
     """
-    Função de inicialização que carrega o modelo VINCIE em uma única GPU.
+    Inicializa e configura o modelo VINCIE.
+    Esta função é chamada uma vez no início da aplicação para carregar o modelo na GPU.
     """
     global MODEL, DEVICE
 
@@ -35,130 +41,265 @@ def setup_model():
         raise RuntimeError("FATAL: Nenhuma GPU compatível com CUDA foi encontrada.")
 
     num_gpus = torch.cuda.device_count()
-    print(f"INFO: Detectadas {num_gpus} GPUs. Usando cuda:0 para a aplicação.")
     if num_gpus == 0:
-        raise RuntimeError("FATAL: Nenhuma GPU foi alocada para este contêiner.")
-
+        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 = "configs/generate.yaml"
-    print(f"INFO: Carregando e resolvendo configuração do modelo de '{config_path}'...")
+    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: Configurando a persistência (modo de inferência)...")
+    # Executa a sequência de inicialização necessária do VINCIE
+    print("INFO: Configurando persistência, modelos e difusão...")
     model_instance.configure_persistence()
-    
-    print("INFO: Configurando os componentes do modelo (DiT, VAE, Text Encoder)...")
-    # O método `configure_models` já move os modelos para o dispositivo CUDA definido
     model_instance.configure_models()
-    
-    print("INFO: Configurando os componentes de difusão...")
     model_instance.configure_diffusion()
     
     if not hasattr(model_instance, 'dit'):
-         raise RuntimeError("FATAL: O modelo 'dit' não foi criado após a configuração.")
+        raise RuntimeError("FATAL: Falha ao inicializar o componente DiT do modelo.")
 
-    # Opcional: garantir que tudo está no dispositivo correto (geralmente já feito)
+    # 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: O modelo VINCIE está pronto na GPU {DEVICE}.")
+    print(f"✅ SUCESSO: Modelo VINCIE pronto para uso na GPU {DEVICE}.")
 
+# --- 3. LÓGICAS DE INFERÊNCIA ---
 
-# --- ETAPA 3: LÓGICA 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 e reutilizável que executa a inferência do VINCIE.
 
-def perform_inference(input_image: str, prompt: str):
-    if MODEL is None:
-        raise gr.Error("O modelo não está carregado. Verifique os logs de inicialização.")
-    if input_image is None or not prompt.strip():
-        raise gr.Error("É necessário fornecer uma imagem de entrada e um prompt de edição.")
-    
-    print(f"INFO: Recebida nova requisição. Prompt: '{prompt}'")
-    turn_1_prompt = [prompt]
-    image_paths = [input_image]
+    Args:
+        prompt_config (DictConfig): Configuração do prompt para o modelo.
+        steps (int): Número de passos de difusão.
+        cfg_scale (float): Escala de orientação (Classifier-Free Guidance).
+        seed (int): Semente para reprodutibilidade.
+        pad_img_placeholder (bool): Se deve formatar o prompt com placeholders <IMG>.
+        resolution (int): Resolução do lado menor da imagem para processamento.
+
+    Returns:
+        Image.Image: A imagem gerada.
+    """
+    # 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:
-        print("INFO: Preparando entradas com `model.prepare_input()`...")
+        OmegaConf.set_readonly(MODEL.config, False)
         
-        prompt_dict = {
-            "index": 0, "img_paths": image_paths, "context": turn_1_prompt,
-        }
-        prompt_config = OmegaConf.create(prompt_dict)
+        # 1. Aplica configurações dinâmicas
+        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)
 
+        # 2. Prepara as entradas
         text_pos, condition, noise, _, _ = MODEL.prepare_input(
-            prompt=prompt_config,
-            repeat_idx=0,
-            device=DEVICE # Usa o dispositivo único definido
+            prompt=prompt_config, repeat_idx=0, device=DEVICE
         )
         
+        # 3. Executa a inferência
         with torch.no_grad():
-            print("INFO: Executando `model.inference()`...")
             samples = MODEL.inference(
                 noises=[noise],
                 conditions=[condition],
                 texts_pos=[text_pos],
                 texts_neg=[MODEL.config.generation.negative_prompt],
+                cfg_scale=cfg_scale
             )
-            print("INFO: Inferência concluída.")
             
         if not samples:
-            raise RuntimeError("A inferência não retornou nenhum resultado.")
+            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 = output_tensor.clip(-1, 1).add(1).div(2).mul(255).byte()
-        output_image = output_image.permute(1, 2, 0).cpu().numpy()
+        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)
+        MODEL.config.diffusion.timesteps.sampling.steps = original_config_state["steps"]
+        MODEL.config.generation.seed = original_config_state["seed"]
+        MODEL.config.generation.pad_img_placehoder = original_config_state["pad"]
+        MODEL.config.generation.resolution = original_config_state["resolution"]
+        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.")
+    
+    _print_params("Edição Simples", prompt=prompt, aspect_ratio=aspect_ratio, resolution=resolution, steps=steps, cfg_scale=cfg_scale, seed=seed)
+    
+    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.")
+        
+    _print_params("Edição em Múltiplos Turnos", prompts=prompts, steps=steps, cfg_scale=cfg_scale, seed=seed)
+    
+    output_images_with_paths = []
+    
+    for i, prompt in enumerate(progress.tqdm(prompts, desc="Processando turnos")):
+        print(f"--- Turno {i+1}/{len(prompts)}: {prompt} ---")
         
-        print("✅ SUCESSO: Imagem processada e retornada para a UI.")
-        return output_image
+        image_paths = [input_image] + [path for _, path in output_images_with_paths]
+        context_prompts = prompts[:i+1]
         
-    except Exception as e:
-        print(f"ERRO: Falha durante a inferência: {e}")
-        import traceback
-        traceback.print_exc()
-        raise gr.Error(f"Ocorreu um erro inesperado. Detalhes: {str(e)}")
+        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.")
+        
+    _print_params("Composição de Conceitos", prompt=prompt, num_images=len(image_paths))
+
+    prefix_prompts = [f"<IMG{i+1}>: " for i in range(len(image_paths) - 1)]
+    all_prompts = prefix_prompts + [prompt]
+
+    prompt_config = OmegaConf.create({
+        "index": 0, "img_paths": image_paths, "context": all_prompts, "aspect_ratio": "1:1"
+    })
+    
+    return _execute_vincie_logic(prompt_config, steps=50, cfg_scale=7.5, seed=1, pad_img_placeholder=False, resolution=512)
+
+
+def _print_params(mode: str, **kwargs: Any):
+    """Função auxiliar para logging formatado das requisições."""
+    log_message = f"\n{'='*50}\nINFO: Nova requisição - Modo: {mode}\n"
+    for key, value in kwargs.items():
+        log_message += f"  - {key.replace('_', ' ').title()}: {value}\n"
+    log_message += f"{'='*50}\n"
+    print(log_message)
 
 
-# --- ETAPA 4 e 5: UI e Lançamento (sem alterações) ---
+# --- 4. CONSTRUÇÃO DA INTERFACE GRadio ---
 
 def create_ui():
-    with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="VINCIE Image Editor") as demo:
-        # ... (código da UI idêntico)
-        gr.Markdown("# 🖼️ **VINCIE: Editor de Imagens por Contexto**\nFaça o upload de uma imagem e descreva a edição que você deseja realizar.")
-        with gr.Row():
-            with gr.Column(scale=1):
-                image_input = gr.Image(type="filepath", label="Imagem de Entrada")
-                prompt_input = gr.Textbox(lines=3, label="O que você quer mudar?", placeholder="Ex: 'mude o fundo para uma praia ensolarada'")
-                submit_button = gr.Button("✨ Gerar Edição", variant="primary")
-            with gr.Column(scale=1):
-                image_output = gr.Image(label="Resultado da Edição", interactive=False, height=512)
-        gr.Examples(
-            examples=[
-                ["/app/VINCIE/assets/woman_pineapple.png", "Adicione uma coroa na cabeça da mulher."],
-                ["/app/VINCIE/assets/da_vinci.png", "Remova o chapéu da cabeça de Leonardo e mude seu casaco para um de cor azul."],
-                ["/app/VINCIE/assets/dog1.png", "Transforme o cachorro em um filhote de leão."],
-            ],
-            inputs=[image_input, prompt_input],
-            outputs=image_output,
-            fn=perform_inference,
-            cache_examples=False,
-        )
-        submit_button.click(fn=perform_inference, inputs=[image_input, prompt_input], outputs=[image_output])
+    """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.")
+
+        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():
+                    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, height=512)
+                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")
+                        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=1):
+                        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>, <IMG1> e <IMG2>, crie um retrato de uma família com o pai de <IMG0>, a mãe de <IMG1> e o cachorro de <IMG2> em um parque. Saída <IMG6>:")
+                concept_button = gr.Button("Compor Imagem", variant="primary")
+                concept_img_out = gr.Image(label="Resultado da Composição", interactive=False, height=512)
+                gr.Examples(
+                    [[
+                        "Baseado em <IMG0>, <IMG1> e <IMG2>, crie um retrato de uma família com o pai de <IMG0>, a mãe de <IMG1> e o cachorro de <IMG2> em um parque. Saída <IMG6>:",
+                        "/app/VINCIE/assets/father.png", "/app/VINCIE/assets/mother.png", "/app/VINCIE/assets/dog1.png"
+                    ]],
+                    [concept_prompt] + concept_inputs
+                )
+
+        # 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__":
     setup_model()
     ui = create_ui()
-    server_name = os.environ.get("GRADIO_SERVER_NAME", "127.0.0.1")
+    
+    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:
-        print("INFO: Fila de requisições (queue) ativada.")
         ui.queue()
+
     ui.launch(server_name=server_name, server_port=server_port)