Update train_lora.py

train_lora.py

@@ -9,22 +9,28 @@ from torchvision import transforms
 from PIL import Image
 import glob
 
+
 def main(args):
+    # Inicializa o Accelerator
     accelerator = Accelerator(
-        mixed_precision="fp16" if args.mixed_precision else None,
-        gradient_accumulation_steps=1
+        mixed_precision="fp16" if args.mixed_precision else None
     )
 
-    # Carrega pipeline
-    print("Carregando modelo base...")
-    pipe = StableDiffusionPipeline.from_pretrained(
-        args.model_name,
-        torch_dtype=torch.float16 if args.mixed_precision else torch.float32
-    )
+    print(f"🚀 Carregando modelo: {args.model_name}")
+    try:
+        pipe = StableDiffusionPipeline.from_pretrained(
+            args.model_name,
+            torch_dtype=torch.float16 if args.mixed_precision else torch.float32
+        )
+    except Exception as e:
+        print(f"❌ Falha ao carregar modelo: {e}")
+        return
+
+    # Extrai componentes
+    unet = pipe.unet
     tokenizer = pipe.tokenizer
     text_encoder = pipe.text_encoder
     vae = pipe.vae
-    unet = pipe.unet
     noise_scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
 
     # Configura LoRA
@@ -36,9 +42,9 @@ def main(args):
         bias="none"
     )
     unet = get_peft_model(unet, lora_config)
-    unet.print_trainable_parameters()
+    unet.print_trainable_parameters()  # Mostra % de parâmetros treináveis
 
-    # Transformações
+    # Transformações de imagem
     transform = transforms.Compose([
         transforms.Resize(512),
         transforms.CenterCrop(512),
@@ -46,9 +52,16 @@ def main(args):
         transforms.Normalize([0.5], [0.5]),
     ])
 
-    # Carrega imagens e legendas
-    image_paths = sorted(glob.glob(os.path.join(args.dataset_dir, "*.*")))
-    image_paths = [p for p in image_paths if p.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.webp'))]
+    # === Carrega dataset ===
+    image_paths = []
+    for ext in ["*.jpg", "*.jpeg", "*.png", "*.bmp", "*.webp"]:
+        image_paths.extend(glob.glob(os.path.join(args.dataset_dir, ext)))
+    
+    if len(image_paths) == 0:
+        print("❌ Nenhuma imagem encontrada no diretório!")
+        return
+
+    print(f"✅ {len(image_paths)} imagens encontradas. Carregando legendas...")
 
     captions = []
     valid_images = []
@@ -56,89 +69,106 @@ def main(args):
         txt_path = os.path.splitext(img_path)[0] + ".txt"
         if os.path.exists(txt_path):
             with open(txt_path, "r", encoding="utf-8") as f:
-                captions.append(f.read().strip())
+                caption = f.read().strip()
         else:
-            captions.append("person")
+            caption = "person"
+        captions.append(caption)
         valid_images.append(img_path)
 
-    if len(valid_images) == 0:
-        print("❌ Nenhuma imagem encontrada!")
-        return
-
-    print(f"✅ {len(valid_images)} imagens carregadas")
-
+    # Dataset PyTorch
     class SimpleDataset(torch.utils.data.Dataset):
-        def __init__(self, image_paths, captions, transform):
-            self.image_paths = image_paths
-            self.captions = captions
+        def __init__(self, image_list, caption_list, transform):
+            self.images = image_list
+            self.captions = caption_list
             self.transform = transform
 
         def __len__(self):
-            return len(self.image_paths)
+            return len(self.images)
 
         def __getitem__(self, idx):
-            image = Image.open(self.image_paths[idx]).convert("RGB")
+            image = Image.open(self.images[idx]).convert("RGB")
             image = self.transform(image)
-            caption = self.captions[idx]
-            return {"pixel_values": image, "input_ids": caption}
+            return {"pixel_values": image, "input_ids": self.captions[idx]}
 
     dataset = SimpleDataset(valid_images, captions, transform)
-    dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=True
+    )
 
     # Otimizador
     optimizer = torch.optim.AdamW(unet.parameters(), lr=args.learning_rate)
-    lr_scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer)
 
-    unet, optimizer, dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, dataloader, lr_scheduler)
+    # Prepara com Accelerator
+    unet, optimizer, dataloader = accelerator.prepare(unet, optimizer, dataloader)
+
+    # Coloca VAE e Text Encoder em modo de avaliação (só UNET é treinado)
+    vae.eval()
+    text_encoder.eval()
 
     # Treinamento
     unet.train()
-    global_step = 0
+    step = 0
     for epoch in range(args.num_epochs):
         for batch in dataloader:
             with accelerator.accumulate(unet):
-                pixel_values = batch["pixel_values"].to(accelerator.device)
+                # Gera latents
+                pixel_values = batch["pixel_values"]
                 latents = vae.encode(pixel_values).latent_dist.sample() * 0.18215
 
+                # Adiciona ruído
                 noise = torch.randn_like(latents)
                 bsz = latents.shape[0]
                 timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (bsz,), device=latents.device)
                 noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
 
-                encoder_hidden_states = text_encoder(tokenizer(
+                # Codifica texto
+                inputs = tokenizer(
                     batch["input_ids"],
-                    padding="max_length",
                     max_length=77,
+                    padding="max_length",
                     truncation=True,
                     return_tensors="pt"
-                ).input_ids.to(latents.device))[0]
+                ).to(latents.device)
+                encoder_hidden_states = text_encoder(**inputs)[0]
 
+                # Predição de ruído
                 noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
                 loss = torch.nn.functional.mse_loss(noise_pred, noise)
 
+                # Backpropagation
                 accelerator.backward(loss)
                 optimizer.step()
-                lr_scheduler.step()
                 optimizer.zero_grad()
-                global_step += 1
+                step += 1
 
-    # Salva modelo
+        print(f"Epoch {epoch+1}/{args.num_epochs} - Loss: {loss.item():.4f}")
+
+    # Salva modelo LoRA
     accelerator.wait_for_everyone()
     if accelerator.is_main_process:
+        output_dir = args.output_dir
         unwrapped_unet = accelerator.unwrap_model(unet)
-        unwrapped_unet.save_pretrained(args.output_dir)
-        print(f"✅ Modelo salvo em {args.output_dir}")
+        unwrapped_unet.save_pretrained(output_dir)
+        print(f"✅ Modelo LoRA salvo em: {output_dir}")
+
+        # Opcional: salva também como safetensors
+        from peft import save_model
+        save_model(unwrapped_unet, output_dir)
+
 
 if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--model_name", type=str, default="runwayml/stable-diffusion-v1-5")
-    parser.add_argument("--dataset_dir", type=str, required=True)
-    parser.add_argument("--output_dir", type=str, default="lora-output")
-    parser.add_argument("--lora_rank", type=int, default=4)
-    parser.add_argument("--lora_alpha", type=int, default=32)
-    parser.add_argument("--learning_rate", type=float, default=1e-4)
-    parser.add_argument("--num_epochs", type=int, default=10)
-    parser.add_argument("--batch_size", type=int, default=1)
-    parser.add_argument("--mixed_precision", action="store_true")
+    parser = argparse.ArgumentParser(description="Treina um modelo LoRA para Stable Diffusion")
+    parser.add_argument("--model_name", type=str, default="runwayml/stable-diffusion-v1-5", help="Modelo base do HF")
+    parser.add_argument("--dataset_dir", type=str, required=True, help="Pasta com imagens e .txt")
+    parser.add_argument("--output_dir", type=str, default="lora-output", help="Onde salvar o LoRA")
+    parser.add_argument("--lora_rank", type=int, default=4, help="Rank LoRA (4-64)")
+    parser.add_argument("--lora_alpha", type=int, default=32, help="Alpha LoRA")
+    parser.add_argument("--learning_rate", type=float, default=1e-4, help="Taxa de aprendizado")
+    parser.add_argument("--num_epochs", type=int, default=10, help="Número de épocas")
+    parser.add_argument("--batch_size", type=int, default=1, help="Batch size")
+    parser.add_argument("--mixed_precision", action="store_true", help="Usa FP16")
+
     args = parser.parse_args()
     main(args)