Update app.py

app.py

@@ -21,7 +21,7 @@ from diffusers import (
     AutoencoderKL
 )
 from transformers import CLIPTextModel, CLIPTokenizer
-from peft import LoraConfig, get_peft_model, TaskType
+from peft import LoraConfig
 import logging
 
 # Configurar logging
@@ -86,30 +86,6 @@ class LoRAImageTrainer:
             logger.error(f"Erro ao carregar modelo {model_name}: {str(e)}")  
             raise e  
   
-    def create_lora_config(self,   
-                          r: int = 16,   
-                          lora_alpha: int = 32,   
-                          lora_dropout: float = 0.1,  
-                          target_modules: Optional[List[str]] = None) -> LoraConfig:  
-        """Cria configuração LoRA otimizada para modelos de difusão."""  
-          
-        if target_modules is None:  
-            # Módulos padrão para UNet do Stable Diffusion  
-            target_modules = [  
-                "to_k", "to_q", "to_v", "to_out.0",  
-                "proj_in", "proj_out",  
-                "ff.net.0.proj", "ff.net.2"  
-            ]  
-          
-        return LoraConfig(  
-            r=r,  
-            lora_alpha=lora_alpha,  
-            target_modules=target_modules,  
-            lora_dropout=lora_dropout,  
-            bias="none",  
-            task_type=TaskType.CAUSAL_LM,  # ✅ CORREÇÃO PRINCIPAL: DIFFUSION → CAUSAL_LM
-        )  
-  
     def prepare_image_dataset(self, image_files: List[str], captions: List[str], resolution: int = 512) -> List[Dict]:  
         """Prepara dataset de imagens para treinamento."""  
         dataset = []  
@@ -171,7 +147,7 @@ class LoRAImageTrainer:
                          learning_rate: float = 1e-4,
                          batch_size: int = 1,
                          resolution: int = 512) -> None:
-        """TREINAMENTO REAL DE LoRA PARA IMAGENS."""
+        """TREINAMENTO REAL DE LoRA PARA IMAGENS - CORRIGIDO PARA DIFFUSERS + PEFT."""
         
         try:
             # Atualizar status
@@ -192,14 +168,25 @@ class LoRAImageTrainer:
             text_encoder.requires_grad_(False)
             vae.requires_grad_(False)
 
-            # Configurar LoRA no UNet
-            lora_config = self.create_lora_config(r, lora_alpha, lora_dropout)
-            unet_lora = get_peft_model(unet, lora_config)
-            unet_lora.train()
-            unet_lora.to(self.device)
+            # Criar configuração LoRA
+            lora_config = LoraConfig(
+                r=r,
+                lora_alpha=lora_alpha,
+                target_modules=["to_k", "to_q", "to_v", "to_out.0"],
+                lora_dropout=lora_dropout,
+                bias="none"
+            )
+
+            # Aplicar LoRA ao UNet manualmente, sem usar get_peft_model diretamente
+            unet.add_adapter(lora_config, adapter_name="default")
+
+            # Ativar o adaptador
+            unet.set_adapter("default")
+            unet.train()
+            unet.to(self.device)
 
             # Otimizador
-            optimizer = torch.optim.AdamW(unet_lora.parameters(), lr=learning_rate)
+            optimizer = torch.optim.AdamW(unet.parameters(), lr=learning_rate)
 
             # Preparar scheduler para treinamento
             self.training_jobs[job_id]["status"] = "preparing_data"
@@ -249,7 +236,7 @@ class LoRAImageTrainer:
 
                     # Forward pass
                     encoder_hidden_states = text_encoder(input_ids)[0]
-                    noise_pred = unet_lora(noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states).sample
+                    noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states).sample
 
                     # Calcular perda
                     loss = torch.nn.functional.mse_loss(noise_pred, noise)
@@ -274,17 +261,17 @@ class LoRAImageTrainer:
             output_dir = f"./lora_models/{job_id}"
             os.makedirs(output_dir, exist_ok=True)
 
-            # Salvar apenas os pesos LoRA do UNet
-            unet_lora.save_pretrained(output_dir)
+            # Salvar apenas os adaptadores LoRA
+            unet.save_pretrained(output_dir)
 
             # Criar adapter_config.json
             lora_config_dict = {
                 "r": r,
                 "lora_alpha": lora_alpha,
-                "target_modules": lora_config.target_modules,
+                "target_modules": ["to_k", "to_q", "to_v", "to_out.0"],
                 "lora_dropout": lora_dropout,
                 "bias": "none",
-                "task_type": "CAUSAL_LM",  # ✅ CORREÇÃO AQUI TAMBÉM: DIFFUSION → CAUSAL_LM
+                "task_type": "CAUSAL_LM",  # Mantido por compatibilidade, mas não é usado
                 "base_model_name": model_name,
                 "training_info": {
                     "num_epochs": num_epochs,