Create train_lora.py

train_lora.py

@@ -0,0 +1,69 @@
+import os
+import argparse
+from accelerate import Accelerator
+from diffusers import StableDiffusionPipeline, UNet2DConditionModel
+from peft import LoraConfig, get_peft_model
+from transformers import AutoTokenizer, AutoModel
+
+def main(args):
+    accelerator = Accelerator()
+    
+    # Carrega modelo base
+    pipeline = StableDiffusionPipeline.from_pretrained(
+        args.model_name,
+        revision="fp16" if args.mixed_precision else None,
+        torch_dtype=torch.float16 if args.mixed_precision else None
+    )
+    
+    # Configura LoRA
+    unet = pipeline.unet
+    lora_config = LoraConfig(
+        r=args.lora_rank,
+        lora_alpha=args.lora_alpha,
+        target_modules=["to_q", "to_v"],
+        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
+    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)
+    
+    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()
+    
+    # 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}"
+        )
+
+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("--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")
+    args = parser.parse_args()
+    main(args)