Update app.py

app.py

@@ -35,8 +35,8 @@ class LoRAImageTrainer:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
         self.training_jobs = {}  
         self.models_cache = {}  
-        # ✅ Criar pasta para persistência de jobs
         Path("./jobs").mkdir(exist_ok=True)
+        Path("./lora_models").mkdir(exist_ok=True)
           
     def _save_job_state(self, job_id: str):
         """Salva o estado do job em disco."""
@@ -67,7 +67,6 @@ class LoRAImageTrainer:
             return {"error": "Job não encontrado"}
 
     def get_available_models(self) -> List[str]:  
-        """Retorna lista de modelos base disponíveis para treinamento LoRA."""  
         return [  
             "runwayml/stable-diffusion-v1-5",  
             "stabilityai/stable-diffusion-2-1",  
@@ -75,7 +74,6 @@ class LoRAImageTrainer:
         ]  
   
     def load_base_model(self, model_name: str):  
-        """Carrega modelo base de difusão com otimizações para baixo uso de GPU."""  
         try:  
             if model_name in self.models_cache:  
                 return self.models_cache[model_name]  
@@ -163,7 +161,6 @@ class LoRAImageTrainer:
                          batch_size: int = 1,
                          resolution: int = 512) -> None:
         try:
-            # Inicializar job se não existir
             if job_id not in self.training_jobs:
                 self.training_jobs[job_id] = {
                     "id": job_id,
@@ -178,13 +175,11 @@ class LoRAImageTrainer:
                     "completed_at": None
                 }
 
-            # Atualizar status
             self.training_jobs[job_id]["status"] = "loading_model"
             self.training_jobs[job_id]["progress"] = 5
             self.training_jobs[job_id]["logs"].append(f"{datetime.now().strftime('%H:%M:%S')} - Carregando modelo base: {model_name}")
-            self._save_job_state(job_id)  # ✅ SALVAR ESTADO
+            self._save_job_state(job_id)
 
-            # Carregar modelo base
             pipeline = self.load_base_model(model_name)
             unet = pipeline.unet
             text_encoder = pipeline.text_encoder
@@ -196,6 +191,10 @@ class LoRAImageTrainer:
             text_encoder.requires_grad_(False)
             vae.requires_grad_(False)
 
+            # ✅ CORREÇÃO 1: REMOVER ADAPTADOR EXISTENTE
+            if hasattr(unet, "peft_config") and "default" in unet.peft_config:
+                unet.delete_adapter("default")
+
             lora_config = LoraConfig(
                 r=r,
                 lora_alpha=lora_alpha,
@@ -206,14 +205,22 @@ class LoRAImageTrainer:
 
             unet.add_adapter(lora_config, adapter_name="default")
             unet.set_adapter("default")
+
+            # ✅ CORREÇÃO 2: ATIVAR APENAS PARÂMETROS DO LORA
+            unet.requires_grad_(False)
+            for name, param in unet.named_parameters():
+                if "lora_" in name:
+                    param.requires_grad = True
+
             unet.train()
             unet.to(self.device)
 
-            optimizer = torch.optim.AdamW(unet.parameters(), lr=learning_rate)
+            # Otimizador só nos parâmetros que requerem gradiente
+            optimizer = torch.optim.AdamW([p for p in unet.parameters() if p.requires_grad], lr=learning_rate)
 
             self.training_jobs[job_id]["status"] = "preparing_data"
             self.training_jobs[job_id]["progress"] = 20
-            self._save_job_state(job_id)  # ✅ SALVAR ESTADO
+            self._save_job_state(job_id)
 
             def preprocess_image(image):
                 image = np.array(image).astype(np.float32) / 255.0
@@ -226,7 +233,7 @@ class LoRAImageTrainer:
 
             self.training_jobs[job_id]["status"] = "training"
             self.training_jobs[job_id]["logs"].append(f"{datetime.now().strftime('%H:%M:%S')} - Iniciando treinamento real...")
-            self._save_job_state(job_id)  # ✅ SALVAR ESTADO
+            self._save_job_state(job_id)
 
             for epoch in range(num_epochs):
                 for item in dataset:
@@ -265,15 +272,16 @@ class LoRAImageTrainer:
                     if current_step % max(1, len(dataset)//2) == 0:
                         log_msg = f"Época {epoch+1}, Step {current_step} - Loss: {loss.item():.4f}"
                         self.training_jobs[job_id]["logs"].append(f"{datetime.now().strftime('%H:%M:%S')} - {log_msg}")
-                        self._save_job_state(job_id)  # ✅ SALVAR ESTADO (opcional aqui para logs)
+                        self._save_job_state(job_id)
 
             self.training_jobs[job_id]["status"] = "saving"
             self.training_jobs[job_id]["progress"] = 95
-            self._save_job_state(job_id)  # ✅ SALVAR ESTADO
+            self._save_job_state(job_id)
 
             output_dir = f"./lora_models/{job_id}"
             os.makedirs(output_dir, exist_ok=True)
 
+            # ✅ SALVAR APENAS O LORA
             unet.save_pretrained(
                 output_dir,
                 safe_serialization=True,
@@ -312,7 +320,7 @@ Data: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
             self.training_jobs[job_id]["model_path"] = output_dir
             self.training_jobs[job_id]["completed_at"] = datetime.now().isoformat()
             self.training_jobs[job_id]["logs"].append(f"{datetime.now().strftime('%H:%M:%S')} - ✅ Treinamento concluído! LoRA salvo em {output_dir}")
-            self._save_job_state(job_id)  # ✅ SALVAR ESTADO FINAL
+            self._save_job_state(job_id)
 
             logger.info(f"Treinamento LoRA concluído para job {job_id}")
 
@@ -323,7 +331,7 @@ Data: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
                 self.training_jobs[job_id]["status"] = "error"
                 self.training_jobs[job_id]["error"] = error_msg
                 self.training_jobs[job_id]["logs"].append(f"{datetime.now().strftime('%H:%M:%S')} - ❌ {error_msg}")
-                self._save_job_state(job_id)  # ✅ SALVAR ESTADO DE ERRO
+                self._save_job_state(job_id)
 
     def start_training(self,
                       model_name: str,
@@ -347,7 +355,6 @@ Data: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
             "completed_at": None  
         }  
 
-        # ✅ Salvar estado inicial
         self._save_job_state(job_id)
 
         thread = threading.Thread(  
@@ -512,14 +519,7 @@ def create_gradio_interface():
     return interface
 
 if __name__ == "__main__":
-    # ✅ Criar diretórios necessários
-    os.makedirs("./lora_models", exist_ok=True)
-    os.makedirs("./jobs", exist_ok=True)  # Pasta para persistência de jobs
-
-    # Configurar interface  
     interface = create_gradio_interface()  
-  
-    # Lançar aplicação  
     interface.launch(  
         server_name="0.0.0.0",  
         server_port=7860,