Update train.py

train.py

@@ -1,68 +1,73 @@
-# train.py
-
 import os
 import torch
 from huggingface_hub import snapshot_download
-from peft import LoraConfig, get_peft_model
-
-# 1️⃣ Pick your scheduler class
 from diffusers import (
     StableDiffusionPipeline,
     DPMSolverMultistepScheduler,
-    UNet2DConditionModel,
     AutoencoderKL,
+    UNet2DConditionModel,
 )
 from transformers import CLIPTextModel, CLIPTokenizer
+from peft import LoraConfig, get_peft_model
 
-# ─── 1) CONFIG ────────────────────────────────────────────────────────────────
+# ─── CONFIG ───────────────────────────────────────────────────────────────────
 
-DATA_DIR   = os.getenv("DATA_DIR", "./data")
-MODEL_DIR  = os.getenv("MODEL_DIR", "./hidream-model")
-OUTPUT_DIR = os.getenv("OUTPUT_DIR", "./lora-trained")
+DATA_DIR      = os.getenv("DATA_DIR", "./data")
+MODEL_CACHE   = os.getenv("MODEL_DIR", "./hidream-model")
+OUTPUT_DIR    = os.getenv("OUTPUT_DIR", "./lora-trained")
+REPO_ID       = "HiDream-ai/HiDream-I1-Dev"
 
-# ─── 2) DOWNLOAD OR VERIFY BASE MODEL ──────────────────────────────────────────
+# ─── STEP 1: ENSURE YOU HAVE A COMPLETE SNAPSHOT WITH CONFIGS ─────────────────
 
-if not os.path.isdir(MODEL_DIR):
-    MODEL_DIR = snapshot_download(
-        repo_id="HiDream-ai/HiDream-I1-Dev",
-        local_dir=MODEL_DIR
-    )
+print(f"📥 Downloading full model snapshot to {MODEL_CACHE}")
+MODEL_ROOT = snapshot_download(
+    repo_id=REPO_ID,
+    local_dir=MODEL_CACHE,
+    local_dir_use_symlinks=False,    # force a copy so config.json ends up there
+)
 
-# ─── 3) LOAD EACH PIPELINE COMPONENT ──────────────────────────────────────────
+# ─── STEP 2: LOAD SCHEDULER ────────────────────────────────────────────────────
 
-# 3a) Scheduler
+print("🔄 Loading scheduler")
 scheduler = DPMSolverMultistepScheduler.from_pretrained(
-    MODEL_DIR,
-    subfolder="scheduler"
+    MODEL_ROOT,
+    subfolder="scheduler",
 )
 
-# 3b) VAE
+# ─── STEP 3: LOAD VAE ──────────────────────────────────────────────────────────
+
+print("🔄 Loading VAE")
 vae = AutoencoderKL.from_pretrained(
-    MODEL_DIR,
+    MODEL_ROOT,
     subfolder="vae",
-    torch_dtype=torch.float16
+    torch_dtype=torch.float16,
 ).to("cuda")
 
-# 3c) Text encoder + tokenizer
+# ─── STEP 4: LOAD TEXT ENCODER + TOKENIZER ─────────────────────────────────────
+
+print("🔄 Loading text encoder + tokenizer")
 text_encoder = CLIPTextModel.from_pretrained(
-    MODEL_DIR,
+    MODEL_ROOT,
     subfolder="text_encoder",
-    torch_dtype=torch.float16
+    torch_dtype=torch.float16,
 ).to("cuda")
-tokenizer    = CLIPTokenizer.from_pretrained(
-    MODEL_DIR,
-    subfolder="tokenizer"
+tokenizer = CLIPTokenizer.from_pretrained(
+    MODEL_ROOT,
+    subfolder="tokenizer",
 )
 
-# 3d) U‑Net
+# ─── STEP 5: LOAD U‑NET ───────────────────────────────────────────────────────
+
+print("🔄 Loading U‑Net")
 unet = UNet2DConditionModel.from_pretrained(
-    MODEL_DIR,
+    MODEL_ROOT,
     subfolder="unet",
-    torch_dtype=torch.float16
+    torch_dtype=torch.float16,
 ).to("cuda")
 
-# ─── 4) BUILD THE PIPELINE ────────────────────────────────────────────────────
+# ─── STEP 6: BUILD THE PIPELINE ───────────────────────────────────────────────
 
+print("🌟 Building StableDiffusionPipeline")
 pipe = StableDiffusionPipeline(
     vae=vae,
     text_encoder=text_encoder,
@@ -71,8 +76,9 @@ pipe = StableDiffusionPipeline(
     scheduler=scheduler,
 ).to("cuda")
 
-# ─── 5) APPLY LORA ────────────────────────────────────────────────────────────
+# ─── STEP 7: APPLY LORA ADAPTER ───────────────────────────────────────────────
 
+print("🧠 Applying LoRA adapter")
 lora_config = LoraConfig(
     r=16,
     lora_alpha=16,
@@ -81,15 +87,15 @@ lora_config = LoraConfig(
 )
 pipe.unet = get_peft_model(pipe.unet, lora_config)
 
-# ─── 6) TRAINING LOOP (SIMULATED) ─────────────────────────────────────────────
+# ─── STEP 8: YOUR TRAINING LOOP (SIMULATED) ────────────────────────────────────
 
-print(f"📂 Data at {DATA_DIR}")
+print(f"📂 Loading dataset from: {DATA_DIR}")
 for step in range(100):
-    # … your real data loading + optimizer here …
+    # ←– here’s where you’d load your images, run forward/backward, optimizer, etc.
     print(f"Training step {step+1}/100")
 
-# ─── 7) SAVE THE FINE‑TUNED LO‑RA ─────────────────────────────────────────────
+# ─── STEP 9: SAVE THE FINE‑TUNED LO‑RA WEIGHTS ───────────────────────────────
 
 os.makedirs(OUTPUT_DIR, exist_ok=True)
 pipe.save_pretrained(OUTPUT_DIR)
-print("✅ Done! Saved to", OUTPUT_DIR)
+print("✅ Training complete. Saved to", OUTPUT_DIR)