Update train.py

train.py

@@ -5,97 +5,64 @@ from diffusers import (
     StableDiffusionPipeline,
     DPMSolverMultistepScheduler,
     AutoencoderKL,
-    UNet2DConditionModel,
+    UNet2DConditionModel
 )
 from transformers import CLIPTextModel, CLIPTokenizer
 from peft import LoraConfig, get_peft_model
 
-# ─── CONFIG ───────────────────────────────────────────────────────────────────
+MODEL_ID = "black-forest-labs/FLUX.1-dev"
+dataset_path = "/workspace/data"
+output_dir   = "/workspace/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"
+# 1) grab the model locally
+print("📥 Downloading Flux‑Dev model…")
+model_path = snapshot_download(MODEL_ID, local_dir="./fluxdev-model")
 
-# ─── STEP 1: ENSURE YOU HAVE A COMPLETE SNAPSHOT WITH CONFIGS ─────────────────
-
-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
-)
-
-# ─── STEP 2: LOAD SCHEDULER ────────────────────────────────────────────────────
-
-print("🔄 Loading scheduler")
+# 2) load each piece with its correct subfolder
+print("🔄 Loading scheduler…")
 scheduler = DPMSolverMultistepScheduler.from_pretrained(
-    MODEL_ROOT,
-    subfolder="scheduler",
+    model_path, subfolder="scheduler"
 )
 
-# ─── STEP 3: LOAD VAE ──────────────────────────────────────────────────────────
-
-print("🔄 Loading VAE")
+print("🔄 Loading VAE…")
 vae = AutoencoderKL.from_pretrained(
-    MODEL_ROOT,
-    subfolder="vae",
-    torch_dtype=torch.float16,
-).to("cuda")
-
-# ─── STEP 4: LOAD TEXT ENCODER + TOKENIZER ─────────────────────────────────────
+    model_path, subfolder="vae", torch_dtype=torch.float16
+)
 
-print("🔄 Loading text encoder + tokenizer")
+print("🔄 Loading text encoder + tokenizer…")
 text_encoder = CLIPTextModel.from_pretrained(
-    MODEL_ROOT,
-    subfolder="text_encoder",
-    torch_dtype=torch.float16,
-).to("cuda")
+    model_path, subfolder="text_encoder", torch_dtype=torch.float16
+)
 tokenizer = CLIPTokenizer.from_pretrained(
-    MODEL_ROOT,
-    subfolder="tokenizer",
+    model_path, subfolder="tokenizer"
 )
 
-# ─── STEP 5: LOAD U‑NET ───────────────────────────────────────────────────────
-
-print("🔄 Loading U‑Net")
+print("🔄 Loading U‑Net…")
 unet = UNet2DConditionModel.from_pretrained(
-    MODEL_ROOT,
-    subfolder="unet",
-    torch_dtype=torch.float16,
-).to("cuda")
-
-# ─── STEP 6: BUILD THE PIPELINE ───────────────────────────────────────────────
+    model_path, subfolder="unet", torch_dtype=torch.float16
+)
 
-print("🌟 Building StableDiffusionPipeline")
+# 3) assemble the pipeline
+print("🛠  Assembling pipeline…")
 pipe = StableDiffusionPipeline(
     vae=vae,
     text_encoder=text_encoder,
     tokenizer=tokenizer,
     unet=unet,
-    scheduler=scheduler,
+    scheduler=scheduler
 ).to("cuda")
 
-# ─── STEP 7: APPLY LORA ADAPTER ───────────────────────────────────────────────
-
-print("🧠 Applying LoRA adapter")
-lora_config = LoraConfig(
-    r=16,
-    lora_alpha=16,
-    bias="none",
-    task_type="CAUSAL_LM",
-)
+# 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)
 
-# ─── STEP 8: YOUR TRAINING LOOP (SIMULATED) ────────────────────────────────────
-
-print(f"📂 Loading dataset from: {DATA_DIR}")
+# 5) your training loop (or dummy loop for illustration)
+print("🚀 Starting fine‑tuning…")
 for step in range(100):
-    # ←– here’s where you’d load your images, run forward/backward, optimizer, etc.
     print(f"Training step {step+1}/100")
+    # …insert your actual data‑loader and loss/backprop here…
 
-# ─── STEP 9: SAVE THE FINE‑TUNED LO‑RA WEIGHTS ───────────────────────────────
-
-os.makedirs(OUTPUT_DIR, exist_ok=True)
-pipe.save_pretrained(OUTPUT_DIR)
-print("✅ Training complete. Saved to", OUTPUT_DIR)
+os.makedirs(output_dir, exist_ok=True)
+pipe.save_pretrained(output_dir)
+print("✅ Done. LoRA weights in", output_dir)