Upload sprite_lora_resume_v7.py with huggingface_hub

sprite_lora_resume_v7.py

@@ -0,0 +1,330 @@
+# /// script
+# requires-python = ">=3.10"
+# dependencies = [
+#     "torch>=2.0.0",
+#     "torchvision>=0.15.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.
+Uses standard FluxPipeline from diffusers.
+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
+from huggingface_hub import (
+    hf_hub_download, 
+    snapshot_download, 
+    create_repo, 
+    upload_folder, 
+    login,
+    HfApi
+)
+from diffusers import FluxPipeline
+from transformers import CLIPTokenizer, T5EncoderModel
+from peft import LoraConfig, get_peft_model, set_peft_model_state_dict, get_peft_model_state_dict
+from safetensors.torch import load_file, save_file
+from accelerate import Accelerator
+
+# Configuration
+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 train(token):
+    print("="*70)
+    print("🚀 FLUX.2-klein-4B LoRA Training Script v7")
+    print("="*70)
+    print(f"\n  Base model: {BASE_MODEL}")
+    print(f"  Dataset:    {DATASET_REPO}")
+    print(f"  Output:     {OUTPUT_REPO}")
+    print(f"  Steps:      1000")
+    print(f"  LoRA:       rank=64, alpha=128")
+    print("="*70)
+    
+    # Authenticate
+    print("\n🔑 Authenticating...")
+    login(token=token, add_to_git_credential=False)
+    print("✅ Authenticated")
+    
+    # Download checkpoint
+    print("\n📥 Downloading checkpoint...")
+    os.makedirs("./checkpoint_step500", exist_ok=True)
+    checkpoint_path = hf_hub_download(
+        repo_id=CHECKPOINT_REPO,
+        filename="pytorch_lora_weights.safetensors",
+        repo_type="model",
+        local_dir="./checkpoint_step500",
+        token=token
+    )
+    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",
+        token=token
+    )
+    image_files = list(Path(dataset_path).rglob("*.png"))
+    print(f"✅ Dataset: {len(image_files)} images")
+    
+    # Setup accelerator
+    print("\n⚙️  Setting up accelerator...")
+    accelerator = Accelerator(
+        gradient_accumulation_steps=4,
+        mixed_precision="bf16"
+    )
+    device = accelerator.device
+    
+    print(f"GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU'}")
+    
+    # Load pipeline
+    print(f"\n📥 Loading {BASE_MODEL}...")
+    pipe = FluxPipeline.from_pretrained(
+        BASE_MODEL,
+        torch_dtype=torch.bfloat16,
+        token=token
+    )
+    
+    # Enable CPU offloading for memory efficiency
+    print("💾 Enabling CPU offloading...")
+    pipe.enable_model_cpu_offload()
+    
+    # Apply LoRA
+    print("\n🔧 Applying LoRA (rank=64, alpha=128)...")
+    
+    # Get target modules for FLUX transformer
+    target_modules = []
+    for i in range(19):  # FLUX.2-klein has 19 transformer blocks
+        target_modules.extend([
+            f"transformer_blocks.{i}.attn.to_q",
+            f"transformer_blocks.{i}.attn.to_k",
+            f"transformer_blocks.{i}.attn.to_v",
+            f"transformer_blocks.{i}.attn.to_out.0",
+        ])
+    
+    lora_config = LoraConfig(
+        r=64,
+        lora_alpha=128,
+        lora_dropout=0.0,
+        target_modules=target_modules,
+        use_rslora=True
+    )
+    
+    pipe.transformer = get_peft_model(pipe.transformer, lora_config)
+    pipe.transformer.print_trainable_parameters()
+    
+    # Load checkpoint
+    print("\n🔄 Loading checkpoint from step 500...")
+    state_dict = load_file(checkpoint_path)
+    set_peft_model_state_dict(pipe.transformer, state_dict)
+    print("✅ Checkpoint loaded")
+    
+    global_step = 500
+    
+    # Create output dir
+    output_dir = Path("./output")
+    output_dir.mkdir(exist_ok=True)
+    
+    # Create output repo
+    print(f"\n📤 Creating output repo: {OUTPUT_REPO}")
+    create_repo(OUTPUT_REPO, exist_ok=True, repo_type="model", token=token)
+    
+    # Setup optimizer
+    print("\n⚙️  Setting up optimizer (AdamW 8-bit)...")
+    trainable_params = [p for p in pipe.transformer.parameters() if p.requires_grad]
+    
+    import bitsandbytes as bnb
+    optimizer = bnb.optim.AdamW8bit(
+        trainable_params,
+        lr=1e-4,
+        betas=(0.9, 0.999),
+        weight_decay=0.01
+    )
+    
+    # Simple dataset
+    print("\n📂 Loading dataset...")
+    
+    class SimpleDataset(torch.utils.data.Dataset):
+        def __init__(self, data_dir, resolution=512):
+            self.data_dir = Path(data_dir)
+            self.resolution = resolution
+            self.image_files = sorted(list(self.data_dir.rglob("*.png")))
+            
+        def __len__(self):
+            return len(self.image_files)
+        
+        def __getitem__(self, idx):
+            img_path = self.image_files[idx]
+            caption_path = img_path.with_suffix(".txt")
+            
+            image = Image.open(img_path).convert("RGB")
+            image = image.resize((self.resolution, self.resolution))
+            image = torch.from_numpy(np.array(image)).permute(2, 0, 1).float() / 255.0
+            image = image * 2.0 - 1.0
+            
+            caption = ""
+            if caption_path.exists():
+                caption = caption_path.read_text().strip()
+            
+            return {"images": image, "captions": caption}
+    
+    import numpy as np
+    dataset = SimpleDataset("./training_data/images")
+    dataloader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=1,
+        shuffle=True,
+        num_workers=0
+    )
+    
+    print(f"✅ Dataset loaded: {len(dataset)} images")
+    
+    # Prepare with accelerator
+    pipe.transformer, optimizer, dataloader = accelerator.prepare(
+        pipe.transformer, optimizer, dataloader
+    )
+    
+    # Training loop
+    print("\n" + "="*70)
+    print("🏋️  Starting Training")
+    print("="*70)
+    print(f"Resuming from step {global_step} to step 1000")
+    print(f"Steps remaining: {1000 - global_step}")
+    
+    pipe.transformer.train()
+    
+    progress_bar = tqdm(total=1000, initial=global_step, desc="Training")
+    
+    while global_step < 1000:
+        for batch in dataloader:
+            with accelerator.accumulate(pipe.transformer):
+                # Get batch data
+                images = batch["images"].to(device)
+                captions = batch["captions"]
+                
+                # Add trigger word
+                captions = [f"pixel art sprite, {c}" for c in captions]
+                
+                # Training step (simplified flow matching)
+                with torch.no_grad():
+                    latents = pipe.vae.encode(images).latent_dist.sample()
+                    noise = torch.randn_like(latents)
+                    timesteps = torch.rand(latents.shape[0], device=device) * 1000
+                    
+                    # Flow matching
+                    sigmas = timesteps.view(-1, 1, 1, 1) / 1000
+                    noisy_latents = (1 - sigmas) * latents + sigmas * noise
+                    target = noise - latents
+                
+                # Get text embeddings
+                with torch.no_grad():
+                    prompt_embeds = pipe.encode_prompt(captions)[0]
+                
+                # Predict
+                model_output = pipe.transformer(
+                    hidden_states=noisy_latents,
+                    timestep=timesteps,
+                    encoder_hidden_states=prompt_embeds,
+                    return_dict=False
+                )[0]
+                
+                # Loss
+                loss = F.mse_loss(model_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
+                progress_bar.update(1)
+                progress_bar.set_postfix({"loss": loss.item()})
+                
+                # Save checkpoint every 500 steps
+                if global_step % 500 == 0:
+                    print(f"\n💾 Saving checkpoint at step {global_step}...")
+                    save_dir = output_dir / f"step_{global_step}"
+                    save_dir.mkdir(exist_ok=True)
+                    
+                    unwrapped = accelerator.unwrap_model(pipe.transformer)
+                    save_file(
+                        get_peft_model_state_dict(unwrapped),
+                        save_dir / "pytorch_lora_weights.safetensors"
+                    )
+                    
+                    # Push to hub
+                    upload_folder(
+                        folder_path=save_dir,
+                        repo_id=OUTPUT_REPO,
+                        repo_type="model",
+                        token=token
+                    )
+                    print(f"✅ Checkpoint pushed to {OUTPUT_REPO}")
+                
+                if global_step >= 1000:
+                    break
+    
+    progress_bar.close()
+    
+    # Final save
+    print("\n💾 Saving final checkpoint...")
+    save_dir = output_dir / "final"
+    save_dir.mkdir(exist_ok=True)
+    
+    unwrapped = accelerator.unwrap_model(pipe.transformer)
+    save_file(
+        get_peft_model_state_dict(unwrapped),
+        save_dir / "pytorch_lora_weights.safetensors"
+    )
+    
+    upload_folder(
+        folder_path=save_dir,
+        repo_id=OUTPUT_REPO,
+        repo_type="model",
+        token=token
+    )
+    
+    print("\n" + "="*70)
+    print("✅ Training Complete!")
+    print("="*70)
+    print(f"\n📤 Model saved to: {OUTPUT_REPO}")
+    print(f"   https://huggingface.co/{OUTPUT_REPO}")
+
+def main():
+    # Get token from environment
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        print("❌ HF_TOKEN not found in environment!")
+        sys.exit(1)
+    
+    print(f"Using token: {token[:7]}...")
+    train(token)
+
+if __name__ == "__main__":
+    main()