Upload sprite_lora_final.py with huggingface_hub

sprite_lora_final.py

@@ -0,0 +1,242 @@
+# /// 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",
+#     "huggingface-hub>=0.20.0",
+#     "safetensors>=0.4.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
+import torch.nn.functional as F
+from pathlib import Path
+from tqdm import tqdm
+from PIL import Image
+import numpy as np
+
+# Get token
+token = os.environ.get("HF_TOKEN")
+if not token or token == "$HF_TOKEN":
+    print("ERROR: HF_TOKEN not set")
+    sys.exit(1)
+
+os.environ["HF_TOKEN"] = token
+
+# Import after setting token
+from huggingface_hub import login, hf_hub_download, snapshot_download, create_repo, upload_file
+from diffusers import FluxPipeline
+from peft import LoraConfig, get_peft_model, set_peft_model_state_dict
+from safetensors.torch import load_file, save_file
+from accelerate import Accelerator
+
+CHECKPOINT_REPO = "Limbicnation/sprite-lora-checkpoint-step500"
+DATASET_REPO = "Limbicnation/sprite-lora-training-data"
+OUTPUT_REPO = "Limbicnation/pixel-art-lora"
+BASE_MODEL = "black-forest-labs/FLUX.2-klein-4B"
+
+def main():
+    print("="*70)
+    print("🚀 FLUX.2-klein-4B LoRA Training - Final")
+    print("="*70)
+    print(f"Base model: {BASE_MODEL}")
+    print(f"Output: {OUTPUT_REPO}")
+    print(f"Resume: Step 500 -> 1000")
+    
+    # Login
+    print("\n🔑 Authenticating...")
+    login(token=token, add_to_git_credential=False)
+    print("✅ Authenticated")
+    
+    # Download checkpoint
+    print("\n📥 Downloading checkpoint...")
+    os.makedirs("checkpoint", exist_ok=True)
+    hf_hub_download(
+        repo_id=CHECKPOINT_REPO,
+        filename="pytorch_lora_weights.safetensors",
+        repo_type="model",
+        local_dir="checkpoint",
+        token=token
+    )
+    print("✅ Checkpoint downloaded")
+    
+    # Download dataset
+    print("\n📥 Downloading dataset...")
+    snapshot_download(
+        repo_id=DATASET_REPO,
+        repo_type="dataset",
+        local_dir="data",
+        token=token
+    )
+    image_files = list(Path("data").rglob("*.png"))
+    print(f"✅ Dataset: {len(image_files)} images")
+    
+    # Setup accelerator
+    accelerator = Accelerator(gradient_accumulation_steps=4, mixed_precision="bf16")
+    device = accelerator.device
+    print(f"\n⚙️ Device: {device}")
+    
+    # Load model
+    print(f"\n📥 Loading {BASE_MODEL}...")
+    pipe = FluxPipeline.from_pretrained(
+        BASE_MODEL,
+        torch_dtype=torch.bfloat16,
+        token=token
+    )
+    pipe.enable_model_cpu_offload()
+    print("✅ Model loaded")
+    
+    # Apply LoRA
+    print("\n🔧 Applying LoRA (rank=64, alpha=128)...")
+    target_modules = []
+    for i in range(19):
+        target_modules.extend([
+            f"transformer_blocks.{i}.attn.to_q",
+            f"transformer_blocks.{i}.attn.to_k",
+            f"transformer_blocks.{i}.attn.to_v",
+        ])
+    
+    lora_config = LoraConfig(r=64, lora_alpha=128, target_modules=target_modules, use_rslora=True)
+    pipe.transformer = get_peft_model(pipe.transformer, lora_config)
+    
+    # Load checkpoint
+    print("\n🔄 Loading checkpoint...")
+    state_dict = load_file("checkpoint/pytorch_lora_weights.safetensors")
+    set_peft_model_state_dict(pipe.transformer, state_dict)
+    print("✅ Checkpoint loaded, resuming from step 500")
+    
+    global_step = 500
+    
+    # Create output repo
+    print(f"\n📤 Creating output repo...")
+    create_repo(OUTPUT_REPO, exist_ok=True, repo_type="model", token=token)
+    
+    # Setup optimizer
+    trainable = [p for p in pipe.transformer.parameters() if p.requires_grad]
+    import bitsandbytes as bnb
+    optimizer = bnb.optim.AdamW8bit(trainable, lr=1e-4)
+    
+    # Dataset
+    class Dataset(torch.utils.data.Dataset):
+        def __init__(self, root, res=512):
+            self.imgs = sorted(list(Path(root).rglob("*.png")))
+            self.res = res
+        def __len__(self): return len(self.imgs)
+        def __getitem__(self, idx):
+            img = Image.open(self.imgs[idx]).convert("RGB").resize((self.res, self.res))
+            img = torch.from_numpy(np.array(img)).permute(2,0,1).float()/255.0 * 2 - 1
+            txt = self.imgs[idx].with_suffix(".txt")
+            cap = txt.read_text().strip() if txt.exists() else ""
+            return {"images": img, "captions": cap}
+    
+    dataset = Dataset("data/images")
+    dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True)
+    print(f"✅ Dataset ready: {len(dataset)} images")
+    
+    # Prepare
+    pipe.transformer, optimizer, dataloader = accelerator.prepare(
+        pipe.transformer, optimizer, dataloader
+    )
+    
+    # Training
+    print("\n" + "="*70)
+    print("🏋️ Training: Step 500 -> 1000")
+    print("="*70)
+    
+    pipe.transformer.train()
+    pbar = tqdm(total=1000, initial=global_step, desc="Training")
+    
+    while global_step < 1000:
+        for batch in dataloader:
+            with accelerator.accumulate(pipe.transformer):
+                imgs = batch["images"].to(device)
+                caps = [f"pixel art sprite, {c}" for c in batch["captions"]]
+                
+                with torch.no_grad():
+                    latents = pipe.vae.encode(imgs).latent_dist.sample()
+                    noise = torch.randn_like(latents)
+                    t = torch.rand(latents.shape[0], device=device) * 1000
+                    sigmas = t.view(-1,1,1,1) / 1000
+                    noisy = (1-sigmas)*latents + sigmas*noise
+                    target = noise - latents
+                
+                with torch.no_grad():
+                    prompt_embeds = pipe.encode_prompt(caps)[0]
+                
+                output = pipe.transformer(
+                    hidden_states=noisy,
+                    timestep=t,
+                    encoder_hidden_states=prompt_embeds,
+                    return_dict=False
+                )[0]
+                
+                loss = F.mse_loss(output.float(), target.float())
+                accelerator.backward(loss)
+                
+                if accelerator.sync_gradients:
+                    accelerator.clip_grad_norm_(pipe.transformer.parameters(), 1.0)
+                
+                optimizer.step()
+                optimizer.zero_grad()
+            
+            if accelerator.sync_gradients:
+                global_step += 1
+                pbar.update(1)
+                pbar.set_postfix({"loss": f"{loss.item():.4f}"})
+                
+                if global_step % 500 == 0:
+                    print(f"\n💾 Saving checkpoint at step {global_step}...")
+                    os.makedirs(f"output/step_{global_step}", exist_ok=True)
+                    save_file(
+                        get_peft_model_state_dict(accelerator.unwrap_model(pipe.transformer)),
+                        f"output/step_{global_step}/pytorch_lora_weights.safetensors"
+                    )
+                    upload_file(
+                        path_or_fileobj=f"output/step_{global_step}/pytorch_lora_weights.safetensors",
+                        path_in_repo=f"step_{global_step}/pytorch_lora_weights.safetensors",
+                        repo_id=OUTPUT_REPO,
+                        repo_type="model",
+                        token=token
+                    )
+                    print("✅ Checkpoint saved")
+                
+                if global_step >= 1000:
+                    break
+    
+    pbar.close()
+    
+    # Final save
+    print("\n💾 Saving final model...")
+    os.makedirs("output/final", exist_ok=True)
+    save_file(
+        get_peft_model_state_dict(accelerator.unwrap_model(pipe.transformer)),
+        "output/final/pytorch_lora_weights.safetensors"
+    )
+    upload_file(
+        path_or_fileobj="output/final/pytorch_lora_weights.safetensors",
+        path_in_repo="pytorch_lora_weights.safetensors",
+        repo_id=OUTPUT_REPO,
+        repo_type="model",
+        token=token
+    )
+    
+    print("\n" + "="*70)
+    print("✅ Training Complete!")
+    print("="*70)
+    print(f"\n📤 Model: https://huggingface.co/{OUTPUT_REPO}")
+
+if __name__ == "__main__":
+    main()