Update train.py

train.py

@@ -1,111 +1,74 @@
+# train.py
+
 import os
-import json
-import inspect
-from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
-from peft import LoraConfig, get_peft_model
 import torch
 from huggingface_hub import snapshot_download
+from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
+from peft import LoraConfig, get_peft_model
+
+# ── 1) Configuration ───────────────────────────────────────────────────────────
+
+# Where you put your images + prompts
+DATA_DIR   = os.getenv("DATA_DIR", "./data")
+
+# Where your base model lives (downloaded or cached)
+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 ────────────────────────────────────────
 
-# ─── 1. Read hyperparameters & mode ───────────────────────────────────────────
-model_id     = os.environ.get("BASE_MODEL", "HiDream-ai/HiDream-I1-Dev")
-trigger_word = os.environ.get("TRIGGER_WORD", "default-style")
-num_steps    = int(os.environ.get("NUM_STEPS", 100))
-lora_r       = int(os.environ.get("LORA_R", 16))
-lora_alpha   = int(os.environ.get("LORA_ALPHA", 16))
-LOCAL        = os.environ.get("LOCAL_TRAIN", "").lower() in ("1", "true")
-
-# ─── 2. Set up directories ────────────────────────────────────────────────────
-if LOCAL:
-    DATA_DIR    = os.path.join(os.getcwd(), "data")
-    OUTPUT_DIR  = os.path.join(os.getcwd(), "lora-trained")
-    LOCAL_MODEL = os.path.join(os.getcwd(), "hidream-model")
-    os.makedirs(DATA_DIR,   exist_ok=True)
-    os.makedirs(OUTPUT_DIR, exist_ok=True)
-else:
-    DATA_DIR   = "/tmp/data"
-    OUTPUT_DIR = "/tmp/lora-trained"
-    CACHE_DIR  = "/tmp/hidream-model"
-    os.makedirs(DATA_DIR,   exist_ok=True)
-    os.makedirs(OUTPUT_DIR, exist_ok=True)
-    os.makedirs(CACHE_DIR,  exist_ok=True)
-
-print(f"📂 Dataset directory: {DATA_DIR}", flush=True)
-print(f"📥 Preparing base model: {model_id}", flush=True)
-
-# ─── 3. Resolve model path ────────────────────────────────────────────────────
-def get_model_path():
-    # If local and predownloaded model exists, use it
-    if LOCAL and os.path.isdir(LOCAL_MODEL) and os.path.isfile(os.path.join(LOCAL_MODEL, "config.json")):
-        print(f"✅ Using local model at: {LOCAL_MODEL}", flush=True)
-        return LOCAL_MODEL
-    # Otherwise download (to ~/.cache on local, or /tmp on Spaces)
-    download_kwargs = {} if LOCAL else {"local_dir": CACHE_DIR}
-    path = snapshot_download(model_id, **download_kwargs)
-    print(f"✅ Downloaded model to: {path}", flush=True)
-    return path
-
-model_path = get_model_path()
-
-# ─── 4. Patch model_index.json to remove unsupported scheduler ────────────────
-mi_file = os.path.join(model_path, "model_index.json")
-if os.path.isfile(mi_file):
-    with open(mi_file, "r") as f:
-        mi = json.load(f)
-    if "pipeline" in mi and "scheduler" in mi["pipeline"]:
-        print("🔧 Removing 'scheduler' entry from model_index.json", flush=True)
-        mi["pipeline"].pop("scheduler", None)
-        with open(mi_file, "w") as f:
-            json.dump(mi, f, indent=2)
-
-# ─── 5. Load & filter scheduler_config.json ──────────────────────────────────
-sched_cfg_path = os.path.join(model_path, "scheduler", "scheduler_config.json")
-filtered_cfg = {}
-if os.path.isfile(sched_cfg_path):
-    with open(sched_cfg_path, "r") as f:
-        raw_cfg = json.load(f)
-    sig = inspect.signature(DPMSolverMultistepScheduler.__init__)
-    valid_keys = set(sig.parameters.keys()) - {"self", "args", "kwargs"}
-    filtered_cfg = {k: v for k, v in raw_cfg.items() if k in valid_keys}
-    dropped = set(raw_cfg) - set(filtered_cfg)
-    if dropped:
-        print(f"⚠️ Dropped unsupported scheduler keys: {dropped}", flush=True)
-    try:
-        scheduler = DPMSolverMultistepScheduler(**filtered_cfg)
-        print("✅ Instantiated DPMSolverMultistepScheduler from config", flush=True)
-    except Exception as e:
-        print(f"❌ Failed to init scheduler from config ({e}), using defaults", flush=True)
-        scheduler = DPMSolverMultistepScheduler()
-else:
-    print("⚠️ No scheduler_config.json found; using default DPMSolverMultistepScheduler", flush=True)
-    scheduler = DPMSolverMultistepScheduler()
-
-# ─── 6. Load the Stable Diffusion pipeline ───────────────────────────��────────
-print(f"🔧 Loading pipeline from: {model_path}", flush=True)
+print(f"📂 Loading dataset from: {DATA_DIR}")
+print("📥 Fetching or verifying base model: HiDream-ai/HiDream-I1-Dev")
+
+# 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 ─────────────────────────────────────────────
+
+# 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")
+scheduler = DPMSolverMultistepScheduler.from_pretrained(
+    pretrained_model_name_or_path=MODEL_DIR,
+    subfolder="scheduler"
+)
+
+# ── 4) Build the Stable Diffusion pipeline ────────────────────────────────────
+
+print("🔧 Creating StableDiffusionPipeline with custom scheduler")
 pipe = StableDiffusionPipeline.from_pretrained(
-    model_path,
+    pretrained_model_name_or_path=MODEL_DIR,
+    scheduler=scheduler,
     torch_dtype=torch.float16,
-    scheduler=scheduler
 ).to("cuda")
 
-# ─── 7. Apply LoRA adapters ───────────────────────────────────────────────────
-print(f"🧠 Applying LoRA config (r={lora_r}, α={lora_alpha})", flush=True)
+# ── 5) Apply PEFT LoRA adapters ───────────────────────────────────────────────
+
+print("🧠 Configuring LoRA adapter on U‑Net")
 lora_config = LoraConfig(
-    r=lora_r,
-    lora_alpha=lora_alpha,
+    r=16,
+    lora_alpha=16,
     bias="none",
     task_type="CAUSAL_LM"
 )
 pipe.unet = get_peft_model(pipe.unet, lora_config)
 
-# ─── 8. Training loop stub ─────────────────────────────────────────────────────
-print(f"🚀 Starting fine‑tuning for {num_steps} steps (trigger: {trigger_word})", flush=True)
-for step in range(num_steps):
-    # TODO: replace this stub with your actual training code:
-    #  • Load batches from DATA_DIR
-    #  • Forward/backward pass, optimizer.step(), etc.
-    print(f"🌀 Step {step+1}/{num_steps}", flush=True)
+# ── 6) (Placeholder) Simulate your training loop ─────────────────────────────
+
+print("🚀 Starting fine‑tuning loop (simulated)")
+for step in range(100):
+    # Here you'd load your data, compute loss, do optimizer.step(), etc.
+    print(f" Training step {step+1}/100")
+
+# ── 7) Save your LoRA‑tuned model ────────────────────────────────────────────
 
-# ─── 9. Save the fine‑tuned model ─────────────────────────────────────────────
-print(f"💾 Saving fine‑tuned model to: {OUTPUT_DIR}", flush=True)
+os.makedirs(OUTPUT_DIR, exist_ok=True)
 pipe.save_pretrained(OUTPUT_DIR)
-print("✅ Training complete!", flush=True)
+print("✅ Training complete. Model saved to", OUTPUT_DIR)