Create train.py

train.py

@@ -0,0 +1,86 @@
+import os, torch
+from datasets import load_dataset
+from diffusers import StableDiffusionPipeline, DDPMScheduler
+from peft import LoraConfig, get_peft_model
+from accelerate import Accelerator
+from PIL import Image
+import numpy as np
+
+# --- CONFIG -------------------------------------------------
+RESOLUTION = 512
+BATCH_SIZE = 1
+GRAD_ACC   = 4
+LR         = 1e-4
+MAX_STEPS  = 500
+OUTPUT_DIR = "./lora-out"
+MODEL_ID   = "runwayml/stable-diffusion-v1-5"
+# -----------------------------------------------------------
+
+accelerator = Accelerator()
+
+# 1. Dataset
+dataset = load_dataset("huggan/smithsonian_butterflies_subset", split="train")
+
+def transform(example):
+    image = example["image"].convert("RGB").resize((RESOLUTION, RESOLUTION))
+    return {"pixel_values": np.array(image).astype(np.float32) / 127.5 - 1.0}
+
+dataset = dataset.map(transform, remove_columns=dataset.column_names)
+
+# 2. Load SD pipeline
+pipe = StableDiffusionPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.float32)
+pipe.scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
+pipe.vae.requires_grad_(False)
+pipe.text_encoder.requires_grad_(False)
+pipe.unet.requires_grad_(False)
+
+# 3. Insert LoRA
+lora_config = LoraConfig(
+    r=16, lora_alpha=16, target_modules=["to_k", "to_q", "to_v", "to_out.0"], init_lora_weights="gaussian"
+)
+pipe.unet = get_peft_model(pipe.unet, lora_config)
+
+# 4. Optimizer & dataloader
+from torch.utils.data import DataLoader
+dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True)
+optimizer = torch.optim.AdamW(pipe.unet.parameters(), lr=LR)
+
+pipe.unet, optimizer, dataloader = accelerator.prepare(pipe.unet, optimizer, dataloader)
+
+# 5. Training loop
+pipe.unet.train()
+for step, batch in enumerate(dataloader, 1):
+    latents = pipe.vae.encode(batch["pixel_values"].to(pipe.vae.dtype)).latent_dist.sample()
+    latents = latents * pipe.vae.config.scaling_factor
+    noise = torch.randn_like(latents)
+    timesteps = torch.randint(0, pipe.scheduler.config.num_train_timesteps, (latents.shape[0],))
+    noisy_latents = pipe.scheduler.add_noise(latents, noise, timesteps)
+
+    encoder_hidden_states = pipe.text_encoder(
+        pipe.tokenizer(
+            ["a high-resolution photo of a butterfly"] * latents.shape[0],
+            padding="max_length",
+            max_length=pipe.tokenizer.model_max_length,
+            return_tensors="pt",
+        ).input_ids.to(pipe.text_encoder.device)
+    )[0]
+
+    model_pred = pipe.unet(noisy_latents, timesteps, encoder_hidden_states).sample
+    loss = torch.nn.functional.mse_loss(model_pred, noise)
+    accelerator.backward(loss)
+    if step % GRAD_ACC == 0:
+        optimizer.step()
+        optimizer.zero_grad()
+
+    if accelerator.is_main_process and step % 50 == 0:
+        print(f"Step {step:04d}/{MAX_STEPS} | loss={loss.item():.4f}")
+    if step >= MAX_STEPS:
+        break
+
+# 6. Save LoRA weights
+accelerator.wait_for_everyone()
+unwrapped = accelerator.unwrap_model(pipe.unet)
+unwrapped.save_pretrained(OUTPUT_DIR)
+print("LoRA saved to", OUTPUT_DIR)
+# Move the file to the expected name for Gradio
+os.rename(f"{OUTPUT_DIR}/adapter_model.safetensors", "./pytorch_lora_weights.safetensors")