Update train_lora.py

train_lora.py

@@ -1,58 +1,126 @@
+# train_lora.py
 import os
-import argparse
-from accelerate import Accelerator
-from diffusers import StableDiffusionPipeline, UNet2DConditionModel
+import torch
+from diffusers import StableDiffusionPipeline
 from peft import LoraConfig, get_peft_model
-from transformers import AutoTokenizer, AutoModel
+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()
-    
-    # Carrega modelo base
-    pipeline = StableDiffusionPipeline.from_pretrained(
+    accelerator = Accelerator(mixed_precision="fp16" if args.mixed_precision else None)
+
+    # Carrega pipeline
+    print("Carregando modelo base...")
+    pipe = StableDiffusionPipeline.from_pretrained(
         args.model_name,
-        revision="fp16" if args.mixed_precision else None,
-        torch_dtype=torch.float16 if args.mixed_precision else None
+        torch_dtype=torch.float16 if args.mixed_precision else torch.float32
     )
-    
+    tokenizer = pipe.tokenizer
+    text_encoder = pipe.text_encoder
+    vae = pipe.vae
+    unet = pipe.unet
+
     # Configura LoRA
-    unet = pipeline.unet
     lora_config = LoraConfig(
         r=args.lora_rank,
         lora_alpha=args.lora_alpha,
-        target_modules=["to_q", "to_v"],
+        target_modules=["to_q", "to_v", "to_k", "to_out.0"],
         lora_dropout=0.0,
         bias="none"
     )
     unet = get_peft_model(unet, lora_config)
-    
-    # Prepara dados
-    from datasets import load_dataset
-    dataset = load_dataset("imagefolder", data_dir=args.dataset_dir, split="train")
-    
-    # Treinamento
+
+    # Transformações
+    transform = transforms.Compose([
+        transforms.Resize(512),
+        transforms.CenterCrop(512),
+        transforms.ToTensor(),
+        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'))]
+
+    captions = []
+    valid_images = []
+    for img_path in image_paths:
+        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())
+        else:
+            captions.append("person")
+        valid_images.append(img_path)
+
+    if len(valid_images) == 0:
+        print("❌ Nenhuma imagem encontrada!")
+        return
+
+    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
+            self.captions = captions
+            self.transform = transform
+
+        def __len__(self):
+            return len(self.image_paths)
+
+        def __getitem__(self, idx):
+            image = Image.open(self.image_paths[idx]).convert("RGB")
+            image = self.transform(image)
+            caption = self.captions[idx]
+            return {"pixel_values": image, "input_ids": caption}
+
+    dataset = SimpleDataset(valid_images, captions, transform)
+    dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size, shuffle=True)
+
+    # Otimizador
     optimizer = torch.optim.AdamW(unet.parameters(), lr=args.learning_rate)
-    train_dataloader = torch.utils.data.DataLoader(dataset, batch_size=args.batch_size)
-    
-    unet, optimizer, train_dataloader = accelerator.prepare(unet, optimizer, train_dataloader)
-    
+    lr_scheduler = get_scheduler("constant", optimizer=optimizer, num_warmup_steps=0, num_training_steps=len(dataloader) * args.num_epochs)
+
+    unet, optimizer, dataloader, lr_scheduler = accelerator.prepare(unet, optimizer, dataloader, lr_scheduler)
+
+    # Treinamento
+    unet.train()
     for epoch in range(args.num_epochs):
-        for step, batch in enumerate(train_dataloader):
-            # Lógica de treinamento simplificada (para demonstração)
-            loss = unet(batch["pixel_values"]).sample.mean()
-            accelerator.backward(loss)
-            optimizer.step()
-            optimizer.zero_grad()
-    
+        for batch in dataloader:
+            with accelerator.accumulate(unet):
+                latents = vae.encode(batch["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)
+
+                encoder_hidden_states = text_encoder(tokenizer(
+                    batch["input_ids"],
+                    padding="max_length",
+                    max_length=77,
+                    truncation=True,
+                    return_tensors="pt"
+                ).input_ids.to(latents.device))[0]
+
+                noise_pred = unet(noisy_latents, timesteps, encoder_hidden_states).sample
+                loss = torch.nn.functional.mse_loss(noise_pred, noise)
+
+                accelerator.backward(loss)
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+
     # Salva modelo
-    unet.save_pretrained(args.output_dir)
-    if args.push_to_hub:
-        from huggingface_hub import upload_folder
-        upload_folder(
-            repo_id=args.hub_model_id,
-            folder_path=args.output_dir,
-            commit_message=f"LoRA fine-tuning epoch {epoch}"
-        )
+    accelerator.wait_for_everyone()
+    if accelerator.is_main_process:
+        unwrapped_unet = accelerator.unwrap_model(unet)
+        unwrapped_unet.save_pretrained(args.output_dir)
+        print(f"✅ Modelo salvo em {args.output_dir}")
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
@@ -60,10 +128,10 @@ if __name__ == "__main__":
     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=4)
-    parser.add_argument("--push_to_hub", action="store_true")
-    parser.add_argument("--hub_model_id", type=str, default="my-lora-model")
+    parser.add_argument("--batch_size", type=int, default=1)
+    parser.add_argument("--mixed_precision", action="store_true")
     args = parser.parse_args()
     main(args)