Update train.py

train.py

@@ -11,8 +11,20 @@ from transformers import CLIPTextModel, CLIPTokenizer
 from peft import LoraConfig, get_peft_model
 
 MODEL_ID = "black-forest-labs/FLUX.1-dev"
-dataset_path = "/workspace/data"
-output_dir   = "/workspace/lora-trained"
+
+# download
+model_path = snapshot_download(
+    MODEL_ID,
+    local_dir="./fluxdev-model",
+    use_auth_token=True
+)
+
+# 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…")