Upload sprite_lora_resume_v2.py with huggingface_hub

sprite_lora_resume_v2.py

@@ -0,0 +1,106 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "torch>=2.0.0",
+#     "diffusers>=0.25.0",
+#     "transformers>=4.35.0",
+#     "accelerate>=0.24.0",
+#     "peft>=0.7.0",
+#     "bitsandbytes>=0.41.0",
+#     "huggingface-hub>=0.20.0",
+#     "safetensors>=0.4.0",
+#     "omegaconf>=2.3.0",
+#     "Pillow>=10.0.0",
+#     "numpy>=1.24.0",
+#     "tqdm>=4.66.0",
+# ]
+# ///
+
+"""
+Resume FLUX.2-klein-4B LoRA training from step 500 checkpoint.
+Output: Limbicnation/pixel-art-lora
+"""
+
+import os
+import sys
+import torch
+from pathlib import Path
+from huggingface_hub import hf_hub_download, snapshot_download, create_repo, upload_folder
+
+CHECKPOINT_REPO = "Limbicnation/sprite-lora-checkpoint-step500"
+DATASET_REPO = "Limbicnation/sprite-lora-training-data"
+OUTPUT_REPO = "Limbicnation/pixel-art-lora"
+
+def main():
+    print("="*70)
+    print("🚀 FLUX.2-klein-4B LoRA Training (Resuming from Step 500)")
+    print("="*70)
+    
+    # Download checkpoint
+    print("\n📥 Downloading checkpoint...")
+    checkpoint_path = hf_hub_download(
+        repo_id=CHECKPOINT_REPO,
+        filename="pytorch_lora_weights.safetensors",
+        repo_type="model",
+        local_dir="./checkpoint_step500"
+    )
+    print(f"   ✅ Checkpoint: {checkpoint_path}")
+    
+    # Download dataset
+    print("\n📥 Downloading dataset...")
+    dataset_path = snapshot_download(
+        repo_id=DATASET_REPO,
+        repo_type="dataset",
+        local_dir="./training_data"
+    )
+    image_files = list(Path(dataset_path).rglob("*.png"))
+    print(f"   ✅ Dataset: {len(image_files)} images")
+    
+    # Clone trainer
+    print("\n📥 Setting up trainer...")
+    os.system("git clone https://github.com/Limbicnation/klein-lora-trainer.git 2>/dev/null || true")
+    sys.path.insert(0, "./klein-lora-trainer")
+    
+    from flux2_klein_trainer.config import TrainingConfig, ModelConfig, LoRAConfig, DatasetConfig
+    from flux2_klein_trainer.trainer import KleinLoRATrainer
+    
+    # Config
+    config = TrainingConfig(
+        model=ModelConfig(
+            pretrained_model_name="black-forest-labs/FLUX.2-klein-4B",
+            dtype="bfloat16",
+            enable_cpu_offload=True,
+        ),
+        lora=LoRAConfig(rank=64, alpha=128),
+        dataset=DatasetConfig(
+            data_dir="./training_data/images",
+            caption_ext="txt",
+            resolution=512,
+        ),
+        output_dir="./output",
+        resume_from_checkpoint="./checkpoint_step500",
+        num_train_steps=1000,
+        batch_size=1,
+        gradient_accumulation_steps=4,
+        learning_rate=1e-4,
+        optimizer="adamw_8bit",
+        save_every=500,
+        sample_every=500,
+        trigger_word="pixel art sprite",
+        push_to_hub=True,
+        hub_model_id=OUTPUT_REPO,
+    )
+    
+    print(f"\n📤 Output: {OUTPUT_REPO}")
+    create_repo(OUTPUT_REPO, exist_ok=True, repo_type="model")
+    
+    # Train
+    print("\n🏋️  Starting Training...")
+    trainer = KleinLoRATrainer(config)
+    trainer.train()
+    
+    print("\n✅ Complete!")
+    print(f"📤 Model saved to: {OUTPUT_REPO}")
+
+if __name__ == "__main__":
+    main()