Update app_vince.py

app_vince.py

@@ -3,7 +3,7 @@ import sys
 import torch
 import gradio as gr
 from PIL import Image
-from omegaconf import OmegaConf # Garantir que OmegaConf está importado
+from omegaconf import OmegaConf
 
 # --- ETAPA 1: CONFIGURAÇÃO DO AMBIENTE ---
 
@@ -20,46 +20,55 @@ except ImportError as e:
     print(f"Verifique se o diretório '{VINCIE_DIR}' contém o código-fonte clonado.")
     raise e
 
-# --- ETAPA 2: INICIALIZAÇÃO E PARALELIZAÇÃO DO MODELO ---
-# (Esta seção permanece inalterada)
+# --- ETAPA 2: INICIALIZAÇÃO DO MODELO (Simplificado para uma GPU) ---
 
 MODEL: VINCIEGenerator = None
-DEVICE_PRIMARY: torch.device = None
+DEVICE: torch.device = None
 
 def setup_model():
-    # ... (código de setup_model idêntico à versão anterior)
-    global MODEL, DEVICE_PRIMARY
+    """
+    Função de inicialização que carrega o modelo VINCIE em uma única GPU.
+    """
+    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()
-    print(f"INFO: Detectadas {num_gpus} GPUs.")
+    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.")
-    DEVICE_PRIMARY = torch.device("cuda:0")
-    torch.cuda.set_device(DEVICE_PRIMARY)
+
+    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 = load_config(config_path, [])
+    
     print("INFO: Instanciando VINCIEGenerator...")
     model_instance = VINCIEGenerator(config)
+    
     print("INFO: Configurando a persistência (modo de inferência)...")
     model_instance.configure_persistence()
-    print("INFO: Configurando os componentes do modelo...")
+    
+    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.")
-    model_instance.dit.to(DEVICE_PRIMARY)
-    model_instance.vae.to(DEVICE_PRIMARY)
-    model_instance.text_encoder.to(DEVICE_PRIMARY)
-    print("INFO: Modelos carregados com sucesso na GPU primária.")
-    if num_gpus > 1:
-        print(f"INFO: Ativando torch.nn.DataParallel para distribuir a carga entre {num_gpus} GPUs...")
-        model_instance.dit = torch.nn.DataParallel(model_instance.dit)
-        print("INFO: DataParallel ativado no modelo DiT.")
+
+    # Opcional: garantir que tudo está no dispositivo correto (geralmente já feito)
+    model_instance.dit.to(DEVICE)
+    model_instance.vae.to(DEVICE)
+    model_instance.text_encoder.to(DEVICE)
+    
     MODEL = model_instance
-    print("✅ SUCESSO: O modelo VINCIE está pronto para receber requisições.")
+    print(f"✅ SUCESSO: O modelo VINCIE está pronto na GPU {DEVICE}.")
 
 
 # --- ETAPA 3: LÓGICA DE INFERÊNCIA ---
@@ -69,26 +78,23 @@ def perform_inference(input_image: str, prompt: str):
         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]
+    
     try:
         print("INFO: Preparando entradas com `model.prepare_input()`...")
-
-        # --- CORREÇÃO PRINCIPAL ---
-        # Cria um dicionário Python padrão
+        
         prompt_dict = {
-            "index": 0,
-            "img_paths": image_paths,
-            "context": turn_1_prompt,
+            "index": 0, "img_paths": image_paths, "context": turn_1_prompt,
         }
-        # Converte o dicionário para um OmegaConf DictConfig antes de passar para a função
         prompt_config = OmegaConf.create(prompt_dict)
 
         text_pos, condition, noise, _, _ = MODEL.prepare_input(
-            prompt=prompt_config,  # Passa o objeto DictConfig
+            prompt=prompt_config,
             repeat_idx=0,
-            device=DEVICE_PRIMARY
+            device=DEVICE # Usa o dispositivo único definido
         )
         
         with torch.no_grad():
@@ -100,41 +106,34 @@ def perform_inference(input_image: str, prompt: str):
                 texts_neg=[MODEL.config.generation.negative_prompt],
             )
             print("INFO: Inferência concluída.")
+            
         if not samples:
             raise RuntimeError("A inferência não retornou nenhum resultado.")
+            
         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()
+        
         print("✅ SUCESSO: Imagem processada e retornada para a UI.")
         return output_image
+        
     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 durante o processamento. Detalhes: {str(e)}")
+        raise gr.Error(f"Ocorreu um erro inesperado. Detalhes: {str(e)}")
 
 
-# --- ETAPA 4: CONSTRUÇÃO DA INTERFACE COM GRADIO ---
-# (Esta seção permanece inalterada)
+# --- ETAPA 4 e 5: UI e Lançamento (sem alterações) ---
 
 def create_ui():
-    # ... (código de create_ui idêntico à versão anterior)
     with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue"), title="VINCIE Image Editor") as demo:
-        gr.Markdown(
-            """
-            # 🖼️ **VINCIE: Editor de Imagens por Contexto**
-            Faça o upload de uma imagem e descreva a edição que você deseja realizar.
-            O modelo processará sua requisição usando múltiplas GPUs.
-            """
-        )
+        # ... (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="Exemplo: 'mude o fundo para uma praia ensolarada e adicione óculos de sol na pessoa'"
-                )
+                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)
@@ -149,31 +148,17 @@ def create_ui():
             fn=perform_inference,
             cache_examples=False,
         )
-        submit_button.click(
-            fn=perform_inference,
-            inputs=[image_input, prompt_input],
-            outputs=[image_output]
-        )
+        submit_button.click(fn=perform_inference, inputs=[image_input, prompt_input], outputs=[image_output])
     return demo
 
-
-# --- ETAPA 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_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
-    )
+    ui.launch(server_name=server_name, server_port=server_port)