Update train.py

train.py

@@ -3,72 +3,93 @@
 import os
 import torch
 from huggingface_hub import snapshot_download
-from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
 from peft import LoraConfig, get_peft_model
 
-# ── 1) Configuration ───────────────────────────────────────────────────────────
+# 1️⃣ Pick your scheduler class
+from diffusers import (
+    StableDiffusionPipeline,
+    DPMSolverMultistepScheduler,
+    UNet2DConditionModel,
+    AutoencoderKL,
+)
+from transformers import CLIPTextModel, CLIPTokenizer
 
-# Where you put your images + prompts
-DATA_DIR   = os.getenv("DATA_DIR", "./data")
+# ─── 1) CONFIG ────────────────────────────────────────────────────────────────
 
-# Where your base model lives (downloaded or cached)
+DATA_DIR   = os.getenv("DATA_DIR", "./data")
 MODEL_DIR  = os.getenv("MODEL_DIR", "./hidream-model")
-
-# Where to save your LoRA‑fine‑tuned model
 OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./lora-trained")
 
-# ── 2) Prepare the base model snapshot ────────────────────────────────────────
-
-print(f"📂 Loading dataset from: {DATA_DIR}")
-print("📥 Fetching or verifying base model: HiDream-ai/HiDream-I1-Dev")
+# ─── 2) DOWNLOAD OR VERIFY BASE MODEL ──────────────────────────────────────────
 
-# If you’ve pre‑downloaded into MODEL_DIR, just use it; otherwise pull from HF Hub
 if not os.path.isdir(MODEL_DIR):
     MODEL_DIR = snapshot_download(
         repo_id="HiDream-ai/HiDream-I1-Dev",
         local_dir=MODEL_DIR
     )
 
-# ── 3) Load the scheduler manually ─────────────────────────────────────────────
+# ─── 3) LOAD EACH PIPELINE COMPONENT ──────────────────────────────────────────
 
-# Diffusers’ scheduler config JSON points at FlowMatchLCMScheduler,
-# but your installed version doesn’t have that class. Instead we
-# force‐load DPMSolverMultistepScheduler via `from_pretrained`.
-print(f"🔄 Loading scheduler from: {MODEL_DIR}/scheduler")
+# 3a) Scheduler
 scheduler = DPMSolverMultistepScheduler.from_pretrained(
-    pretrained_model_name_or_path=MODEL_DIR,
+    MODEL_DIR,
     subfolder="scheduler"
 )
 
-# ── 4) Build the Stable Diffusion pipeline ────────────────────────────────────
+# 3b) VAE
+vae = AutoencoderKL.from_pretrained(
+    MODEL_DIR,
+    subfolder="vae",
+    torch_dtype=torch.float16
+).to("cuda")
+
+# 3c) Text encoder + tokenizer
+text_encoder = CLIPTextModel.from_pretrained(
+    MODEL_DIR,
+    subfolder="text_encoder",
+    torch_dtype=torch.float16
+).to("cuda")
+tokenizer    = CLIPTokenizer.from_pretrained(
+    MODEL_DIR,
+    subfolder="tokenizer"
+)
+
+# 3d) U‑Net
+unet = UNet2DConditionModel.from_pretrained(
+    MODEL_DIR,
+    subfolder="unet",
+    torch_dtype=torch.float16
+).to("cuda")
+
+# ─── 4) BUILD THE PIPELINE ────────────────────────────────────────────────────
 
-print("🔧 Creating StableDiffusionPipeline with custom scheduler")
-pipe = StableDiffusionPipeline.from_pretrained(
-    pretrained_model_name_or_path=MODEL_DIR,
+pipe = StableDiffusionPipeline(
+    vae=vae,
+    text_encoder=text_encoder,
+    tokenizer=tokenizer,
+    unet=unet,
     scheduler=scheduler,
-    torch_dtype=torch.float16,
 ).to("cuda")
 
-# ── 5) Apply PEFT LoRA adapters ───────────────────────────────────────────────
+# ─── 5) APPLY LORA ────────────────────────────────────────────────────────────
 
-print("🧠 Configuring LoRA adapter on U‑Net")
 lora_config = LoraConfig(
     r=16,
     lora_alpha=16,
     bias="none",
-    task_type="CAUSAL_LM"
+    task_type="CAUSAL_LM",
 )
 pipe.unet = get_peft_model(pipe.unet, lora_config)
 
-# ── 6) (Placeholder) Simulate your training loop ─────────────────────────────
+# ─── 6) TRAINING LOOP (SIMULATED) ─────────────────────────────────────────────
 
-print("🚀 Starting fine‑tuning loop (simulated)")
+print(f"📂 Data at {DATA_DIR}")
 for step in range(100):
-    # Here you'd load your data, compute loss, do optimizer.step(), etc.
-    print(f" Training step {step+1}/100")
+    # … your real data loading + optimizer here …
+    print(f"Training step {step+1}/100")
 
-# ── 7) Save your LoRA‑tuned model ────────────────────────────────────────────
+# ─── 7) SAVE THE FINE‑TUNED LO‑RA ─────────────────────────────────────────────
 
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 pipe.save_pretrained(OUTPUT_DIR)
-print("✅ Training complete. Model saved to", OUTPUT_DIR)
+print("✅ Done! Saved to", OUTPUT_DIR)