Update app.py

app.py

@@ -35,14 +35,20 @@ class LoRAImageTrainer:
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  
         self.training_jobs = {}  
         self.models_cache = {}  
-        Path("./jobs").mkdir(exist_ok=True)
-        Path("./lora_models").mkdir(exist_ok=True)
+        # ✅ Garantir que as pastas existam no diretório atual
+        os.makedirs("./lora_models", exist_ok=True)
+        os.makedirs("./jobs", exist_ok=True)
+        logger.info("Pastas ./lora_models e ./jobs criadas com sucesso.")
           
     def _save_job_state(self, job_id: str):
         """Salva o estado do job em disco."""
         job_file = Path(f"./jobs/{job_id}.json")
-        with open(job_file, "w") as f:
-            json.dump(self.training_jobs[job_id], f, indent=2, default=str)
+        try:
+            with open(job_file, "w") as f:
+                json.dump(self.training_jobs[job_id], f, indent=2, default=str)
+            logger.info(f"Estado do job {job_id} salvo em disco.")
+        except Exception as e:
+            logger.error(f"Erro ao salvar job {job_id}: {e}")
 
     def _load_job_state(self, job_id: str) -> Optional[Dict]:
         """Carrega o estado do job do disco."""
@@ -50,9 +56,13 @@ class LoRAImageTrainer:
         if job_file.exists():
             try:
                 with open(job_file, "r") as f:
-                    return json.load(f)
+                    loaded_data = json.load(f)
+                    logger.info(f"Estado do job {job_id} carregado do disco.")
+                    return loaded_data
             except Exception as e:
                 logger.error(f"Erro ao carregar job {job_id}: {e}")
+        else:
+            logger.warning(f"Arquivo do job {job_id} não encontrado em disco.")
         return None
 
     def get_training_status(self, job_id: str) -> Dict[str, Any]:
@@ -177,8 +187,10 @@ class LoRAImageTrainer:
 
             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}")
+            log_msg = f"{datetime.now().strftime('%H:%M:%S')} - Carregando modelo base: {model_name}"
+            self.training_jobs[job_id]["logs"].append(log_msg)
             self._save_job_state(job_id)
+            logger.info(log_msg)
 
             pipeline = self.load_base_model(model_name)
             unet = pipeline.unet
@@ -191,9 +203,10 @@ class LoRAImageTrainer:
             text_encoder.requires_grad_(False)
             vae.requires_grad_(False)
 
-            # ✅ CORREÇÃO 1: REMOVER ADAPTADOR EXISTENTE
+            # ✅ Remover adaptador existente
             if hasattr(unet, "peft_config") and "default" in unet.peft_config:
                 unet.delete_adapter("default")
+                logger.info("Adaptador 'default' removido com sucesso.")
 
             lora_config = LoraConfig(
                 r=r,
@@ -206,16 +219,19 @@ class LoRAImageTrainer:
             unet.add_adapter(lora_config, adapter_name="default")
             unet.set_adapter("default")
 
-            # ✅ CORREÇÃO 2: ATIVAR APENAS PARÂMETROS DO LORA
+            # ✅ Ativar apenas parâmetros do LoRA
             unet.requires_grad_(False)
+            trainable_params = 0
             for name, param in unet.named_parameters():
                 if "lora_" in name:
                     param.requires_grad = True
+                    trainable_params += 1
+
+            logger.info(f"Número de parâmetros treináveis (LoRA): {trainable_params}")
 
             unet.train()
             unet.to(self.device)
 
-            # 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"
@@ -232,8 +248,10 @@ class LoRAImageTrainer:
             current_step = 0
 
             self.training_jobs[job_id]["status"] = "training"
-            self.training_jobs[job_id]["logs"].append(f"{datetime.now().strftime('%H:%M:%S')} - Iniciando treinamento real...")
+            log_msg = f"{datetime.now().strftime('%H:%M:%S')} - Iniciando treinamento real..."
+            self.training_jobs[job_id]["logs"].append(log_msg)
             self._save_job_state(job_id)
+            logger.info(log_msg)
 
             for epoch in range(num_epochs):
                 for item in dataset:
@@ -271,8 +289,9 @@ 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.training_jobs[job_id]["logs"].append(log_msg)
                         self._save_job_state(job_id)
+                        logger.info(log_msg)
 
             self.training_jobs[job_id]["status"] = "saving"
             self.training_jobs[job_id]["progress"] = 95
@@ -281,7 +300,6 @@ class LoRAImageTrainer:
             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,
@@ -319,10 +337,10 @@ Data: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
             self.training_jobs[job_id]["progress"] = 100
             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}")
+            log_msg = f"{datetime.now().strftime('%H:%M:%S')} - ✅ Treinamento concluído! LoRA salvo em {output_dir}"
+            self.training_jobs[job_id]["logs"].append(log_msg)
             self._save_job_state(job_id)
-
-            logger.info(f"Treinamento LoRA concluído para job {job_id}")
+            logger.info(log_msg)
 
         except Exception as e:
             error_msg = f"Erro no treinamento: {str(e)}"
@@ -386,7 +404,8 @@ Data: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
                                 "r": config.get("r", "Unknown"),  
                                 "created": datetime.fromtimestamp(model_dir.stat().st_mtime).isoformat()  
                             })  
-                        except:  
+                        except Exception as e:  
+                            logger.error(f"Erro ao ler config de {model_dir.name}: {e}")
                             models.append({  
                                 "id": model_dir.name,  
                                 "path": str(model_dir),  
@@ -519,6 +538,11 @@ def create_gradio_interface():
     return interface
 
 if __name__ == "__main__":
+    # ✅ Garantir que os diretórios existam
+    os.makedirs("./lora_models", exist_ok=True)
+    os.makedirs("./jobs", exist_ok=True)
+    logger.info("Aplicação iniciada. Diretórios verificados.")
+
     interface = create_gradio_interface()  
     interface.launch(  
         server_name="0.0.0.0",