Update train.py

train.py

@@ -1,80 +1,55 @@
 import os
 import torch
-from huggingface_hub import snapshot_download
-from diffusers import (
-    StableDiffusionPipeline,
-    DPMSolverMultistepScheduler,
-    AutoencoderKL,
-    UNet2DConditionModel
+from aitoolkit import (
+    LoRATrainer,
+    StableDiffusionModel,
+    LoRAConfig,
+    ImageTextDataset,
 )
-from transformers import CLIPTextModel, CLIPTokenizer
-from peft import LoraConfig, get_peft_model
 
-MODEL_ID = "black-forest-labs/FLUX.1-dev"
+# 1. Configuration
+MODEL_ID     = "HiDream-ai/HiDream-I1-Dev"    # or your gated FLUX model if you have access
+DATA_DIR     = "/workspace/data"
+OUTPUT_DIR   = "/workspace/lora-trained"
+DEVICE       = "cuda" if torch.cuda.is_available() else "cpu"
 
-# download
-model_path = snapshot_download(
-    MODEL_ID,
-    local_dir="./fluxdev-model",
-    use_auth_token=True
+lora_cfg = LoRAConfig(
+    rank=16,
+    alpha=16,
+    bias="none",
 )
 
-# later loading
-pipe = StableDiffusionPipeline.from_pretrained(
-    model_path,
-    torch_dtype=torch.float16,
-    use_auth_token=True
-).to("cuda")
-
-# 1) grab the model locally
-print("📥 Downloading Flux‑Dev model…")
-model_path = snapshot_download(MODEL_ID, local_dir="./fluxdev-model")
+training_args = {
+    "num_train_steps": 100,
+    "batch_size": 4,
+    "learning_rate": 1e-4,
+    "save_every_n_steps": 50,
+    "output_dir": OUTPUT_DIR,
+}
 
-# 2) load each piece with its correct subfolder
-print("🔄 Loading scheduler…")
-scheduler = DPMSolverMultistepScheduler.from_pretrained(
-    model_path, subfolder="scheduler"
+# 2. Load base diffusion model
+model = StableDiffusionModel.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16,
+    device=DEVICE,
+    use_auth_token=True,   # if it’s a gated repo
 )
 
-print("🔄 Loading VAE…")
-vae = AutoencoderKL.from_pretrained(
-    model_path, subfolder="vae", torch_dtype=torch.float16
-)
+# 3. Prepare your dataset
+# Expects pairs of image files + .txt captions in DATA_DIR
+dataset = ImageTextDataset(data_root=DATA_DIR, image_size=512)
 
-print("🔄 Loading text encoder + tokenizer…")
-text_encoder = CLIPTextModel.from_pretrained(
-    model_path, subfolder="text_encoder", torch_dtype=torch.float16
-)
-tokenizer = CLIPTokenizer.from_pretrained(
-    model_path, subfolder="tokenizer"
-)
+# 4. Hook up the LoRA adapter
+model.apply_lora(lora_cfg)
 
-print("🔄 Loading U‑Net…")
-unet = UNet2DConditionModel.from_pretrained(
-    model_path, subfolder="unet", torch_dtype=torch.float16
+# 5. Create the trainer and kickoff
+trainer = LoRATrainer(
+    model=model,
+    dataset=dataset,
+    args=training_args,
 )
 
-# 3) assemble the pipeline
-print("🛠  Assembling pipeline…")
-pipe = StableDiffusionPipeline(
-    vae=vae,
-    text_encoder=text_encoder,
-    tokenizer=tokenizer,
-    unet=unet,
-    scheduler=scheduler
-).to("cuda")
-
-# 4) apply LoRA
-print("🧠 Applying LoRA…")
-lora_config = LoraConfig(r=16, lora_alpha=16, bias="none", task_type="CAUSAL_LM")
-pipe.unet = get_peft_model(pipe.unet, lora_config)
-
-# 5) your training loop (or dummy loop for illustration)
-print("🚀 Starting fine‑tuning…")
-for step in range(100):
-    print(f"Training step {step+1}/100")
-    # …insert your actual data‑loader and loss/backprop here…
+print("🚀 Starting training with AI‑Toolkit…")
+trainer.train()
 
-os.makedirs(output_dir, exist_ok=True)
-pipe.save_pretrained(output_dir)
-print("✅ Done. LoRA weights in", output_dir)
+print(f"✅ Done! Fine-tuned weights saved to {OUTPUT_DIR}")