Update app.py

app.py

@@ -35,7 +35,37 @@ 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)
           
+    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)
+
+    def _load_job_state(self, job_id: str) -> Optional[Dict]:
+        """Carrega o estado do job do disco."""
+        job_file = Path(f"./jobs/{job_id}.json")
+        if job_file.exists():
+            try:
+                with open(job_file, "r") as f:
+                    return json.load(f)
+            except Exception as e:
+                logger.error(f"Erro ao carregar job {job_id}: {e}")
+        return None
+
+    def get_training_status(self, job_id: str) -> Dict[str, Any]:
+        """Retorna status do treinamento, carregando do disco se necessário."""
+        if job_id in self.training_jobs:
+            return self.training_jobs[job_id]
+        else:
+            loaded_job = self._load_job_state(job_id)
+            if loaded_job:
+                self.training_jobs[job_id] = loaded_job
+                return loaded_job
+            return {"error": "Job não encontrado"}
+
     def get_available_models(self) -> List[str]:  
         """Retorna lista de modelos base disponíveis para treinamento LoRA."""  
         return [  
@@ -52,7 +82,6 @@ class LoRAImageTrainer:
               
             logger.info(f"Carregando modelo base: {model_name}")  
               
-            # Configurações para otimização de memória  
             model_kwargs = {  
                 "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32,  
                 "use_safetensors": True,  
@@ -60,7 +89,6 @@ class LoRAImageTrainer:
                 "safety_checker": None,  
             }  
               
-            # Carregar pipeline completo  
             pipeline = StableDiffusionPipeline.from_pretrained(  
                 model_name,  
                 **model_kwargs  
@@ -73,7 +101,6 @@ class LoRAImageTrainer:
                     pipeline.enable_xformers_memory_efficient_attention()  
                 except:  
                     logger.warning("xformers não disponível")  
-                # ✅ ATIVAÇÃO DO GRADIENT CHECKPOINTING — REDUZ MEMÓRIA EM ATÉ 60%
                 pipeline.unet.enable_gradient_checkpointing()
               
             self.models_cache[model_name] = pipeline  
@@ -84,7 +111,6 @@ class LoRAImageTrainer:
             raise e  
   
     def prepare_image_dataset(self, image_files: List[str], captions: List[str], resolution: int = 512) -> List[Dict]:  
-        """Prepara dataset de imagens para treinamento."""  
         dataset = []  
           
         for img_path, caption in zip(image_files, captions):  
@@ -105,7 +131,6 @@ class LoRAImageTrainer:
         return dataset  
   
     def resize_image(self, image: Image.Image, target_size: int) -> Image.Image:  
-        """Redimensiona imagem mantendo aspect ratio e fazendo crop central."""  
         width, height = image.size  
           
         if width > height:  
@@ -137,14 +162,28 @@ class LoRAImageTrainer:
                          learning_rate: float = 1e-4,
                          batch_size: int = 1,
                          resolution: int = 512) -> None:
-        """TREINAMENTO REAL DE LoRA PARA IMAGENS - VERSÃO FINAL OTIMIZADA."""
-        
         try:
+            # Inicializar job se não existir
+            if job_id not in self.training_jobs:
+                self.training_jobs[job_id] = {
+                    "id": job_id,
+                    "status": "queued",
+                    "progress": 0,
+                    "created_at": datetime.now().isoformat(),
+                    "model_name": model_name,
+                    "num_images": len(dataset),
+                    "logs": [],
+                    "error": None,
+                    "model_path": None,
+                    "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
+
             # Carregar modelo base
             pipeline = self.load_base_model(model_name)
             unet = pipeline.unet
@@ -153,12 +192,10 @@ class LoRAImageTrainer:
             tokenizer = pipeline.tokenizer
             scheduler = pipeline.scheduler
 
-            # Congelar parâmetros
             unet.requires_grad_(False)
             text_encoder.requires_grad_(False)
             vae.requires_grad_(False)
 
-            # Criar configuração LoRA
             lora_config = LoraConfig(
                 r=r,
                 lora_alpha=lora_alpha,
@@ -167,18 +204,16 @@ class LoRAImageTrainer:
                 bias="none"
             )
 
-            # Aplicar LoRA ao UNet
             unet.add_adapter(lora_config, adapter_name="default")
             unet.set_adapter("default")
             unet.train()
             unet.to(self.device)
 
-            # Otimizador
             optimizer = torch.optim.AdamW(unet.parameters(), lr=learning_rate)
 
-            # Preparar dados
             self.training_jobs[job_id]["status"] = "preparing_data"
             self.training_jobs[job_id]["progress"] = 20
+            self._save_job_state(job_id)  # ✅ SALVAR ESTADO
 
             def preprocess_image(image):
                 image = np.array(image).astype(np.float32) / 255.0
@@ -186,74 +221,65 @@ class LoRAImageTrainer:
                 image = torch.from_numpy(image).unsqueeze(0)
                 return image
 
-            # Loop de treinamento
             total_steps = num_epochs * len(dataset)
             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...")
+            self._save_job_state(job_id)  # ✅ SALVAR ESTADO
 
             for epoch in range(num_epochs):
                 for item in dataset:
                     current_step += 1
                     
-                    # Pré-processar imagem
                     image = item["image"]
                     caption = item["caption"]
                     image_tensor = preprocess_image(image).to(self.device)
                     if torch.cuda.is_available():
                         image_tensor = image_tensor.half()
 
-                    # Codificar para latentes
                     with torch.no_grad():
                         latents = vae.encode(image_tensor * 2 - 1).latent_dist.sample() * 0.18215
 
-                    # Tokenizar texto
                     inputs = tokenizer(caption, padding="max_length", max_length=tokenizer.model_max_length, truncation=True, return_tensors="pt")
                     input_ids = inputs.input_ids.to(self.device)
 
-                    # Gerar timesteps
                     timesteps = torch.randint(0, scheduler.config.num_train_timesteps, (1,), device=self.device).long()
                     noise = torch.randn_like(latents)
                     noisy_latents = scheduler.add_noise(latents, noise, timesteps)
 
-                    # Forward pass
                     encoder_hidden_states = text_encoder(input_ids)[0]
                     noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states=encoder_hidden_states).sample
 
-                    # Calcular e propagar perda
                     loss = torch.nn.functional.mse_loss(noise_pred, noise)
                     optimizer.zero_grad()
                     loss.backward()
                     optimizer.step()
 
-                    # ✅ LIMPEZA DE MEMÓRIA A CADA STEP
                     if torch.cuda.is_available():
                         torch.cuda.empty_cache()
 
-                    # Atualizar progresso
                     progress = 30 + int((current_step / total_steps) * 60)
                     self.training_jobs[job_id]["progress"] = min(progress, 90)
 
                     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)
 
-            # ✅ SALVAR APENAS OS ADAPTADORES (NÃO O MODELO INTEIRO)
             self.training_jobs[job_id]["status"] = "saving"
             self.training_jobs[job_id]["progress"] = 95
+            self._save_job_state(job_id)  # ✅ SALVAR ESTADO
 
             output_dir = f"./lora_models/{job_id}"
             os.makedirs(output_dir, exist_ok=True)
 
-            # 👇👇👇 CORREÇÃO FINAL: SALVA SÓ O LORA 👇👇👇
             unet.save_pretrained(
                 output_dir,
                 safe_serialization=True,
                 selected_adapters=["default"]
             )
 
-            # Criar adapter_config.json
             lora_config_dict = {
                 "r": r,
                 "lora_alpha": lora_alpha,
@@ -273,7 +299,6 @@ class LoRAImageTrainer:
             with open(f"{output_dir}/adapter_config.json", "w") as f:
                 json.dump(lora_config_dict, f, indent=2)
 
-            # README
             readme_content = f"""# LoRA Model - {job_id}
 Treinado com sucesso!
 Modelo Base: {model_name}
@@ -282,21 +307,23 @@ Data: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
             with open(f"{output_dir}/README.md", "w") as f:
                 f.write(readme_content)
 
-            # Finalizar
             self.training_jobs[job_id]["status"] = "completed"
             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}")
+            self._save_job_state(job_id)  # ✅ SALVAR ESTADO FINAL
 
             logger.info(f"Treinamento LoRA concluído para job {job_id}")
 
         except Exception as e:
             error_msg = f"Erro no treinamento: {str(e)}"
             logger.error(error_msg)
-            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}")
+            if job_id in self.training_jobs:
+                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
 
     def start_training(self,
                       model_name: str,
@@ -304,6 +331,7 @@ Data: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
                       captions: List[str],
                       **kwargs) -> str:
         job_id = str(uuid.uuid4())  
+
         dataset = self.prepare_image_dataset(image_files, captions, kwargs.get('resolution', 512))  
 
         self.training_jobs[job_id] = {  
@@ -319,6 +347,9 @@ 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(  
             target=self.real_lora_training,  
             args=(job_id, model_name, dataset),  
@@ -329,9 +360,6 @@ Data: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
 
         return job_id
 
-    def get_training_status(self, job_id: str) -> Dict[str, Any]:
-        return self.training_jobs.get(job_id, {"error": "Job não encontrado"})
-
     def list_trained_models(self) -> List[Dict[str, str]]:
         models = []
         lora_models_dir = Path("./lora_models")
@@ -400,7 +428,6 @@ def create_gradio_interface():
                 captions.append(f"{trigger_word.strip() or 'training image'}, high quality photo" if trigger_word.strip() else f"training image {len(captions) + 1}, high quality photo")  
             captions = captions[:len(files)]  
             
-            # ✅ VALORES FIXOS DEFINIDOS AQUI DENTRO
             job_id = trainer.start_training(  
                 model_name=model_name, 
                 image_files=image_files, 
@@ -465,7 +492,6 @@ def create_gradio_interface():
                         learning_rate = gr.Slider(1e-5, 1e-3, 1e-4, step=1e-5, label="Taxa de Aprendizado")  
                 train_button = gr.Button("🚀 Iniciar Treinamento", variant="primary")  
                 train_output = gr.Textbox(label="📊 Resultado")  
-                # ✅ APENAS COMPONENTES GRADIO SÃO PASSADOS
                 train_button.click(start_training_wrapper, [model_dropdown, image_files, captions_text, trigger_word, r, lora_alpha, num_epochs, learning_rate], train_output)  
             
             with gr.TabItem("📊 Status"):  
@@ -486,6 +512,17 @@ 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()  
-    interface.launch(server_name="0.0.0.0", server_port=7860, show_error=True)
+  
+    # Lançar aplicação  
+    interface.launch(  
+        server_name="0.0.0.0",  
+        server_port=7860,  
+        show_error=True,  
+        quiet=False  
+    )