Update train_lora.py

train_lora.py

@@ -1,17 +1,19 @@
 # train_lora.py
 import os
 import torch
-from diffusers import StableDiffusionPipeline
+import argparse
+from diffusers import StableDiffusionPipeline, DDPMScheduler
 from peft import LoraConfig, get_peft_model
-from diffusers.optimization import get_scheduler
 from accelerate import Accelerator
 from torchvision import transforms
 from PIL import Image
-import argparse
 import glob
 
 def main(args):
-    accelerator = Accelerator(mixed_precision="fp16" if args.mixed_precision else None)
+    accelerator = Accelerator(
+        mixed_precision="fp16" if args.mixed_precision else None,
+        gradient_accumulation_steps=1
+    )
 
     # Carrega pipeline
     print("Carregando modelo base...")
@@ -23,6 +25,7 @@ def main(args):
     text_encoder = pipe.text_encoder
     vae = pipe.vae
     unet = pipe.unet
+    noise_scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
 
     # Configura LoRA
     lora_config = LoraConfig(
@@ -33,6 +36,7 @@ def main(args):
         bias="none"
     )
     unet = get_peft_model(unet, lora_config)
+    unet.print_trainable_parameters()
 
     # Transformações
     transform = transforms.Compose([
@@ -63,7 +67,6 @@ def main(args):
 
     print(f"✅ {len(valid_images)} imagens carregadas")
 
-    # Dataset simples
     class SimpleDataset(torch.utils.data.Dataset):
         def __init__(self, image_paths, captions, transform):
             self.image_paths = image_paths
@@ -84,20 +87,23 @@ def main(args):
 
     # Otimizador
     optimizer = torch.optim.AdamW(unet.parameters(), lr=args.learning_rate)
-    lr_scheduler = get_scheduler("constant", optimizer=optimizer, num_warmup_steps=0, num_training_steps=len(dataloader) * args.num_epochs)
+    lr_scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer)
 
     unet, optimizer, dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, dataloader, lr_scheduler)
 
     # Treinamento
     unet.train()
+    global_step = 0
     for epoch in range(args.num_epochs):
         for batch in dataloader:
             with accelerator.accumulate(unet):
-                latents = vae.encode(batch["pixel_values"]).latent_dist.sample() * 0.18215
+                pixel_values = batch["pixel_values"].to(accelerator.device)
+                latents = vae.encode(pixel_values).latent_dist.sample() * 0.18215
+
                 noise = torch.randn_like(latents)
                 bsz = latents.shape[0]
-                timesteps = torch.randint(0, 1000, (bsz,), device=latents.device)
-                noisy_latents = latents + noise * torch.sqrt(timesteps / 1000)
+                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(
                     batch["input_ids"],
@@ -114,6 +120,7 @@ def main(args):
                 optimizer.step()
                 lr_scheduler.step()
                 optimizer.zero_grad()
+                global_step += 1
 
     # Salva modelo
     accelerator.wait_for_everyone()