Spaces:
Running on Zero
Running on Zero
| import os | |
| import sys | |
| import gc | |
| import re | |
| import glob | |
| import json | |
| import uuid | |
| import time | |
| import random | |
| import shutil | |
| import zipfile | |
| import tempfile | |
| import threading | |
| from pathlib import Path | |
| from typing import List, Optional, Dict, Any, Tuple | |
| os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1") | |
| sys.path.insert(0, os.getcwd()) | |
| # Safe import for Hugging Face ZeroGPU (`spaces`) | |
| try: | |
| import spaces | |
| IS_ZEROGPU_SPACE = True | |
| except ImportError: | |
| IS_ZEROGPU_SPACE = False | |
| class _DummySpaces: | |
| def GPU(func=None, **kwargs): | |
| if func is not None and callable(func): | |
| return func | |
| def decorator(f): | |
| return f | |
| return decorator | |
| spaces = _DummySpaces() | |
| import yaml | |
| import torch | |
| import gradio as gr | |
| from PIL import Image | |
| from slugify import slugify | |
| from huggingface_hub import HfApi, hf_hub_download, snapshot_download, whoami | |
| from safetensors.torch import load_file | |
| from toolkit.job import get_job | |
| from toolkit.accelerator import reset_accelerator | |
| from toolkit.metadata import load_metadata_from_safetensors | |
| MAX_IMAGES = 150 | |
| SUPPORTED_IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".bmp") | |
| # ============================================================================== | |
| # Model Presets & Architecture Configurations | |
| # ============================================================================== | |
| MODEL_PRESETS: Dict[str, Dict[str, Any]] = { | |
| "Illustrious XL v0.1 (Diffusers Repo)": { | |
| "name_or_path": "OnomaAIResearch/Illustrious-xl-early-release-v0", | |
| "arch": "sdxl", | |
| "is_xl": True, | |
| "is_flux": False, | |
| "quantize": False, | |
| "quantize_te": False, | |
| "noise_scheduler": "ddpm", | |
| "sampler": "euler_a", | |
| "lr": 1e-4, | |
| "unet_lr": 1e-4, | |
| "text_encoder_lr": 5e-5, | |
| "train_text_encoder": True, | |
| "min_snr_gamma": 5.0, | |
| "resolution": [768, 1024], | |
| "guidance_scale": 6.0, | |
| "sample_steps": 28, | |
| "default_neg": "worst quality, low quality, oldest, blurry, bad anatomy, bad hands, text, watermark, signature", | |
| "shuffle_tokens": True, | |
| "description": "Illustrious XL v0.1 by OnomaAIResearch (SDXL anime/illustration foundation model, Diffusers format).", | |
| }, | |
| "Illustrious XL v0.1 (Single-File .safetensors)": { | |
| "name_or_path": "OnomaAIResearch/Illustrious-xl-early-release-v0/Illustrious-XL-v0.1.safetensors", | |
| "arch": "sdxl", | |
| "is_xl": True, | |
| "is_flux": False, | |
| "quantize": False, | |
| "quantize_te": False, | |
| "noise_scheduler": "ddpm", | |
| "sampler": "euler_a", | |
| "lr": 1e-4, | |
| "unet_lr": 1e-4, | |
| "text_encoder_lr": 5e-5, | |
| "train_text_encoder": True, | |
| "min_snr_gamma": 5.0, | |
| "resolution": [768, 1024], | |
| "guidance_scale": 6.0, | |
| "sample_steps": 28, | |
| "default_neg": "worst quality, low quality, oldest, blurry, bad anatomy, bad hands, text, watermark, signature", | |
| "shuffle_tokens": True, | |
| "description": "Illustrious XL v0.1 loaded directly from the official single-file `.safetensors` checkpoint on Hugging Face Hub.", | |
| }, | |
| "Illustrious XL v0.1 GUIDED (Single-File .safetensors)": { | |
| "name_or_path": "OnomaAIResearch/Illustrious-xl-early-release-v0/Illustrious-XL-v0.1-GUIDED.safetensors", | |
| "arch": "sdxl", | |
| "is_xl": True, | |
| "is_flux": False, | |
| "quantize": False, | |
| "quantize_te": False, | |
| "noise_scheduler": "ddpm", | |
| "sampler": "euler_a", | |
| "lr": 1e-4, | |
| "unet_lr": 1e-4, | |
| "text_encoder_lr": 5e-5, | |
| "train_text_encoder": True, | |
| "min_snr_gamma": 5.0, | |
| "resolution": [768, 1024], | |
| "guidance_scale": 6.0, | |
| "sample_steps": 28, | |
| "default_neg": "worst quality, low quality, oldest, blurry, bad anatomy, bad hands, text, watermark, signature", | |
| "shuffle_tokens": True, | |
| "description": "Illustrious XL v0.1 GUIDED variant loaded from official single-file `.safetensors` on Hugging Face Hub.", | |
| }, | |
| "FLUX.1-dev": { | |
| "name_or_path": "black-forest-labs/FLUX.1-dev", | |
| "arch": "flux", | |
| "is_xl": False, | |
| "is_flux": True, | |
| "quantize": True, | |
| "quantize_te": True, | |
| "noise_scheduler": "flowmatch", | |
| "sampler": "flowmatch", | |
| "lr": 1e-4, | |
| "unet_lr": 1e-4, | |
| "text_encoder_lr": 1e-4, | |
| "train_text_encoder": False, | |
| "min_snr_gamma": None, | |
| "resolution": [512, 768, 1024], | |
| "guidance_scale": 3.5, | |
| "sample_steps": 24, | |
| "default_neg": "", | |
| "shuffle_tokens": False, | |
| "description": "FLUX.1-dev 12B rectified flow transformer (requires accepting gate on HF Hub + HF_TOKEN).", | |
| }, | |
| "FLUX.1-schnell (w/ Training Adapter)": { | |
| "name_or_path": "black-forest-labs/FLUX.1-schnell", | |
| "assistant_lora_path": "ostris/FLUX.1-schnell-training-adapter", | |
| "arch": "flux", | |
| "is_xl": False, | |
| "is_flux": True, | |
| "quantize": True, | |
| "quantize_te": True, | |
| "noise_scheduler": "flowmatch", | |
| "sampler": "flowmatch", | |
| "lr": 1e-4, | |
| "unet_lr": 1e-4, | |
| "text_encoder_lr": 1e-4, | |
| "train_text_encoder": False, | |
| "min_snr_gamma": None, | |
| "resolution": [512, 768, 1024], | |
| "guidance_scale": 1.0, | |
| "sample_steps": 4, | |
| "default_neg": "", | |
| "shuffle_tokens": False, | |
| "description": "Apache-2.0 licensed FLUX.1-schnell paired with Ostris' training adapter.", | |
| }, | |
| "Stable Diffusion XL Base 1.0": { | |
| "name_or_path": "stabilityai/stable-diffusion-xl-base-1.0", | |
| "arch": "sdxl", | |
| "is_xl": True, | |
| "is_flux": False, | |
| "quantize": False, | |
| "quantize_te": False, | |
| "noise_scheduler": "ddpm", | |
| "sampler": "ddpm", | |
| "lr": 1e-4, | |
| "unet_lr": 1e-4, | |
| "text_encoder_lr": 5e-5, | |
| "train_text_encoder": True, | |
| "min_snr_gamma": 5.0, | |
| "resolution": [768, 1024], | |
| "guidance_scale": 6.0, | |
| "sample_steps": 25, | |
| "default_neg": "blurry, low quality, deformed", | |
| "shuffle_tokens": False, | |
| "description": "Standard SDXL 1.0 Base model.", | |
| }, | |
| "Custom Diffusers / HF Hub Model": { | |
| "name_or_path": "", | |
| "arch": "sdxl", | |
| "is_xl": True, | |
| "is_flux": False, | |
| "quantize": False, | |
| "quantize_te": False, | |
| "noise_scheduler": "ddpm", | |
| "sampler": "euler_a", | |
| "lr": 1e-4, | |
| "unet_lr": 1e-4, | |
| "text_encoder_lr": 5e-5, | |
| "train_text_encoder": True, | |
| "min_snr_gamma": 5.0, | |
| "resolution": [768, 1024], | |
| "guidance_scale": 6.0, | |
| "sample_steps": 25, | |
| "default_neg": "worst quality, low quality, blurry", | |
| "shuffle_tokens": False, | |
| "description": "Specify any Hugging Face Diffusers repo ID (`org/model`) or single-file checkpoint (`org/model/weights.safetensors`) and select its architecture.", | |
| }, | |
| } | |
| CUSTOM_ARCHITECTURES = [ | |
| ("SDXL / Illustrious / Pony / NoobAI (sdxl)", "sdxl"), | |
| ("FLUX.1 (flux)", "flux"), | |
| ("SD 1.5 (sd1)", "sd1"), | |
| ("SD 2.1 (sd2)", "sd2"), | |
| ("SD 3 / SD 3.5 (sd3)", "sd3"), | |
| ("Lumina 2 (lumina2)", "lumina2"), | |
| ("PixArt Sigma (pixart_sigma)", "pixart_sigma"), | |
| ("AuraFlow (auraflow)", "auraflow"), | |
| ("Flex.1 (flex1)", "flex1"), | |
| ("flex.2 (flex2)", "flex2"), | |
| ("Chroma (chroma)", "chroma"), | |
| ("Z-Image (zimage)", "zimage"), | |
| ("Z-Image Turbo (zimage:turbo)", "zimage:turbo"), | |
| ("Qwen-Image (qwen_image)", "qwen_image"), | |
| ("Anima (anima)", "anima"), | |
| ("Wan 2.1 (wan21)", "wan21"), | |
| ] | |
| FLOWMATCH_ARCHS = { | |
| "flux", "sd3", "lumina2", "auraflow", "flex1", "flex2", | |
| "chroma", "zimage", "zimage:turbo", "qwen_image", "anima", "wan21" | |
| } | |
| # ============================================================================== | |
| # Helper Functions (CPU & File Management) | |
| # ============================================================================== | |
| def resolve_hf_token(user_token: Optional[str] = None, oauth_token: Optional[Any] = None) -> Optional[str]: | |
| """Resolves the Hugging Face token from UI input, OAuth login, or environment.""" | |
| if user_token and str(user_token).strip(): | |
| token = str(user_token).strip() | |
| os.environ["HF_TOKEN"] = token | |
| return token | |
| if oauth_token is not None: | |
| token_str = getattr(oauth_token, "token", None) or (oauth_token if isinstance(oauth_token, str) else None) | |
| if token_str and str(token_str).strip(): | |
| os.environ["HF_TOKEN"] = str(token_str).strip() | |
| return str(token_str).strip() | |
| env_tok = os.environ.get("HF_TOKEN") | |
| if env_tok and env_tok.strip(): | |
| return env_tok.strip() | |
| return None | |
| def get_hf_username(token: Optional[str] = None) -> Optional[str]: | |
| try: | |
| info = whoami(token=token) | |
| return info.get("name") | |
| except Exception: | |
| return None | |
| def extract_uploaded_files(uploaded_files: Optional[List[Any]]) -> Tuple[List[str], Dict[str, str]]: | |
| """ | |
| Processes uploaded files (including .zip archives, images, and .txt captions). | |
| Returns (sorted_image_paths, {stem: caption_text}). | |
| """ | |
| if not uploaded_files: | |
| return [], {} | |
| extract_dir = os.path.join("tmp", "uploads", str(uuid.uuid4())) | |
| os.makedirs(extract_dir, exist_ok=True) | |
| raw_files: List[str] = [] | |
| for f in uploaded_files: | |
| fpath = f.name if hasattr(f, "name") else str(f) | |
| if fpath.lower().endswith(".zip"): | |
| zip_sub = os.path.join(extract_dir, "zip_" + str(uuid.uuid4())[:8]) | |
| os.makedirs(zip_sub, exist_ok=True) | |
| with zipfile.ZipFile(fpath, "r") as zf: | |
| zf.extractall(zip_sub) | |
| for root, _, files in os.walk(zip_sub): | |
| if "__MACOSX" in root: | |
| continue | |
| for name in files: | |
| if not name.startswith("."): | |
| raw_files.append(os.path.join(root, name)) | |
| else: | |
| raw_files.append(fpath) | |
| image_paths: List[str] = [] | |
| captions_map: Dict[str, str] = {} | |
| for fpath in raw_files: | |
| ext = os.path.splitext(fpath)[1].lower() | |
| stem = os.path.splitext(os.path.basename(fpath))[0] | |
| if ext in SUPPORTED_IMAGE_EXTS: | |
| image_paths.append(fpath) | |
| elif ext == ".txt": | |
| try: | |
| with open(fpath, "r", encoding="utf-8", errors="ignore") as tf: | |
| captions_map[stem] = tf.read().strip() | |
| except Exception: | |
| pass | |
| image_paths.sort(key=lambda p: os.path.basename(p).lower()) | |
| return image_paths, captions_map | |
| def create_dataset_folder( | |
| image_paths: List[str], | |
| captions: List[str], | |
| dataset_dir: Optional[str] = None, | |
| ) -> str: | |
| """ | |
| Copies images into `dataset_dir` and writes BOTH `<image_stem>.txt` caption files | |
| (required by ai-toolkit's CaptionMixin dataloader) and `metadata.jsonl`. | |
| """ | |
| if not dataset_dir: | |
| dataset_dir = os.path.join("datasets", str(uuid.uuid4())) | |
| os.makedirs(dataset_dir, exist_ok=True) | |
| jsonl_path = os.path.join(dataset_dir, "metadata.jsonl") | |
| with open(jsonl_path, "w", encoding="utf-8") as jf: | |
| for idx, img_path in enumerate(image_paths): | |
| if not img_path or not os.path.exists(img_path): | |
| continue | |
| ext = os.path.splitext(img_path)[1].lower() | |
| if ext not in SUPPORTED_IMAGE_EXTS: | |
| ext = ".png" | |
| base_stem = f"img_{idx:04d}_{slugify(os.path.splitext(os.path.basename(img_path))[0])[:32]}" | |
| dest_img_name = f"{base_stem}{ext}" | |
| dest_img_path = os.path.join(dataset_dir, dest_img_name) | |
| shutil.copy2(img_path, dest_img_path) | |
| caption_text = (captions[idx] if idx < len(captions) and captions[idx] else "").strip() | |
| dest_txt_path = os.path.join(dataset_dir, f"{base_stem}.txt") | |
| with open(dest_txt_path, "w", encoding="utf-8") as tf: | |
| tf.write(caption_text) | |
| jf.write(json.dumps({"file_name": dest_img_name, "prompt": caption_text}) + "\n") | |
| return dataset_dir | |
| def get_checkpoint_status(training_folder: str, lora_name: str) -> Tuple[int, Optional[str], List[str]]: | |
| """ | |
| Inspects `training_folder/lora_name` for existing checkpoints and sample images. | |
| Returns (current_step, latest_safetensors_path, list_of_sample_images). | |
| """ | |
| slugged_name = slugify(lora_name) if lora_name else "my_lora" | |
| save_root = os.path.join(training_folder, slugged_name) | |
| if not os.path.exists(save_root) and os.path.exists(os.path.join(training_folder, lora_name)): | |
| save_root = os.path.join(training_folder, lora_name) | |
| slugged_name = lora_name | |
| if not os.path.exists(save_root): | |
| return 0, None, [] | |
| safetensors_files = glob.glob(os.path.join(save_root, f"{slugged_name}*.safetensors")) | |
| safetensors_files = [p for p in safetensors_files if os.path.isfile(p)] | |
| best_step = 0 | |
| latest_ckpt = None | |
| if safetensors_files: | |
| def _sort_key(p: str): | |
| m = re.search(r"_(\d{9})\.safetensors$", os.path.basename(p)) | |
| s = int(m.group(1)) if m else -1 | |
| return (s, os.path.getmtime(p)) | |
| safetensors_files.sort(key=_sort_key) | |
| latest_ckpt = safetensors_files[-1] | |
| # Read step from metadata or filename | |
| try: | |
| meta = load_metadata_from_safetensors(latest_ckpt) | |
| if meta and "training_info" in meta and "step" in meta["training_info"]: | |
| best_step = int(meta["training_info"]["step"]) | |
| except Exception: | |
| pass | |
| if best_step == 0: | |
| m = re.search(r"_(\d{9})\.safetensors$", os.path.basename(latest_ckpt)) | |
| if m: | |
| best_step = int(m.group(1)) | |
| samples_dir = os.path.join(save_root, "samples") | |
| sample_imgs: List[str] = [] | |
| if os.path.isdir(samples_dir): | |
| for ext in ("*.jpg", "*.png", "*.webp"): | |
| sample_imgs.extend(glob.glob(os.path.join(samples_dir, ext))) | |
| sample_imgs.sort(key=os.path.getmtime, reverse=True) | |
| return best_step, latest_ckpt, sample_imgs[:12] | |
| def sync_checkpoint_from_hub( | |
| repo_id: str, | |
| training_folder: str, | |
| lora_name: str, | |
| token: Optional[str] = None, | |
| ) -> str: | |
| """ | |
| Downloads existing checkpoints, optimizer.pt, and ema.pt from a Hugging Face Hub repository | |
| into `training_folder/lora_name` so chunked training can resume even after a Space restart. | |
| """ | |
| slugged_name = slugify(lora_name) if lora_name else "my_lora" | |
| save_root = os.path.join(training_folder, slugged_name) | |
| if not os.path.exists(save_root) and os.path.exists(os.path.join(training_folder, lora_name)): | |
| save_root = os.path.join(training_folder, lora_name) | |
| slugged_name = lora_name | |
| os.makedirs(save_root, exist_ok=True) | |
| snapshot_download( | |
| repo_id=repo_id, | |
| repo_type="model", | |
| local_dir=save_root, | |
| allow_patterns=["*.safetensors", "optimizer.pt", "ema.pt", "learnable_snr.json", "samples/*"], | |
| token=token, | |
| ) | |
| step, ckpt, _ = get_checkpoint_status(training_folder, lora_name) | |
| if ckpt: | |
| return f"Downloaded checkpoint from `{repo_id}` at step **{step}** (`{os.path.basename(ckpt)}`)." | |
| return f"Connected to `{repo_id}`, no existing checkpoints found yet (starting from step 0)." | |
| def recursive_update(d: dict, u: dict) -> dict: | |
| for k, v in u.items(): | |
| if isinstance(v, dict) and v: | |
| d[k] = recursive_update(d.get(k, {}), v) | |
| else: | |
| d[k] = v | |
| return d | |
| def build_training_config( | |
| lora_name: str, | |
| trigger_word: str, | |
| preset_name: str, | |
| custom_model_path: str, | |
| custom_arch: str, | |
| dataset_folder: str, | |
| total_steps: int, | |
| chunk_steps: int, | |
| lr: float, | |
| unet_lr: float, | |
| text_encoder_lr: float, | |
| train_text_encoder: bool, | |
| rank: int, | |
| alpha: int, | |
| batch_size: int, | |
| grad_accum: int, | |
| resolution_choice: List[str], | |
| optimizer_type: str, | |
| lr_scheduler_type: str, | |
| quantize_model: bool, | |
| quantize_te: bool, | |
| low_vram: bool, | |
| use_ema: bool, | |
| cache_latents: bool, | |
| cache_text_embeddings: bool, | |
| shuffle_tokens: bool, | |
| save_every: int, | |
| sample_every: int, | |
| disable_sampling: bool, | |
| sample_prompts: List[str], | |
| negative_prompt: str, | |
| guidance_scale: float, | |
| sample_steps: int, | |
| push_to_hub: bool, | |
| push_every_chunk: bool, | |
| hf_repo_id: str, | |
| hf_private: bool, | |
| use_advanced_yaml: bool, | |
| advanced_yaml_str: str, | |
| training_folder: str = "output", | |
| ) -> Tuple[Dict[str, Any], str]: | |
| """Builds the ai-toolkit job config dictionary and writes it to a YAML file.""" | |
| slugged_name = slugify(lora_name) if lora_name else "my_lora" | |
| preset = MODEL_PRESETS.get(preset_name, MODEL_PRESETS["Illustrious XL v0.1 (Diffusers Repo)"]) | |
| if preset_name == "Custom Diffusers / HF Hub Model": | |
| if not custom_model_path or not custom_model_path.strip(): | |
| raise gr.Error("Please enter a Hugging Face model ID or single-file `.safetensors` path for your Custom Model.") | |
| model_path = custom_model_path.strip() | |
| arch = custom_arch.strip() | |
| is_xl = arch in ("sdxl", "illustrious", "vega", "ssd") | |
| is_flux = arch in ("flux", "flex1") | |
| is_flow = arch in FLOWMATCH_ARCHS | |
| noise_scheduler = "flowmatch" if is_flow else "ddpm" | |
| sampler = "flowmatch" if is_flow else "euler_a" | |
| min_snr_gamma = None if is_flow else 5.0 | |
| else: | |
| model_path = preset["name_or_path"] | |
| if custom_model_path and custom_model_path.strip(): | |
| # Allow overriding path even on presets (e.g. custom Illustrious finetune) | |
| model_path = custom_model_path.strip() | |
| arch = preset["arch"] | |
| is_xl = preset["is_xl"] | |
| is_flux = preset["is_flux"] | |
| noise_scheduler = preset["noise_scheduler"] | |
| sampler = preset["sampler"] | |
| min_snr_gamma = preset.get("min_snr_gamma", None) | |
| resolutions = [int(r) for r in resolution_choice] if resolution_choice else [1024] | |
| # Flux / FlowMatch models generally cannot train text encoder directly with cache_text_embeddings | |
| if is_flux and train_text_encoder: | |
| train_text_encoder = False | |
| if cache_text_embeddings and train_text_encoder: | |
| # Cannot train TE while caching text embeddings | |
| cache_text_embeddings = False | |
| model_cfg: Dict[str, Any] = { | |
| "name_or_path": model_path, | |
| "arch": arch, | |
| "quantize": bool(quantize_model), | |
| "quantize_te": bool(quantize_te), | |
| "low_vram": bool(low_vram), | |
| } | |
| if is_xl: | |
| model_cfg["is_xl"] = True | |
| if is_flux: | |
| model_cfg["is_flux"] = True | |
| if "assistant_lora_path" in preset and preset["assistant_lora_path"]: | |
| model_cfg["assistant_lora_path"] = preset["assistant_lora_path"] | |
| valid_prompts = [p.strip() for p in sample_prompts if p and p.strip()] | |
| do_disable_sampling = bool(disable_sampling) or len(valid_prompts) == 0 | |
| train_cfg: Dict[str, Any] = { | |
| "batch_size": int(batch_size), | |
| "steps": int(total_steps), | |
| "max_steps_per_run": int(chunk_steps) if chunk_steps and int(chunk_steps) > 0 else None, | |
| "gradient_accumulation_steps": int(grad_accum), | |
| "train_unet": True, | |
| "train_text_encoder": bool(train_text_encoder), | |
| "gradient_checkpointing": True, | |
| "noise_scheduler": noise_scheduler, | |
| "optimizer": optimizer_type, | |
| "lr_scheduler": lr_scheduler_type, | |
| "lr": float(lr), | |
| "unet_lr": float(unet_lr), | |
| "text_encoder_lr": float(text_encoder_lr), | |
| "skip_first_sample": True, | |
| "disable_sampling": do_disable_sampling, | |
| "cache_text_embeddings": bool(cache_text_embeddings), | |
| "ema_config": { | |
| "use_ema": bool(use_ema), | |
| "ema_decay": 0.99, | |
| }, | |
| "dtype": "bf16", | |
| } | |
| if min_snr_gamma is not None: | |
| train_cfg["min_snr_gamma"] = float(min_snr_gamma) | |
| process_dict: Dict[str, Any] = { | |
| "type": "sd_trainer", | |
| "training_folder": training_folder, | |
| "device": "cuda:0", | |
| "network": { | |
| "type": "lora", | |
| "linear": int(rank), | |
| "linear_alpha": int(alpha), | |
| }, | |
| "save": { | |
| "dtype": "float16", | |
| "save_every": int(save_every), | |
| "max_step_saves_to_keep": 4, | |
| "push_to_hub": bool(push_to_hub), | |
| "push_every_chunk": bool(push_every_chunk), | |
| "hf_repo_id": hf_repo_id if push_to_hub else None, | |
| "hf_private": bool(hf_private), | |
| }, | |
| "datasets": [ | |
| { | |
| "folder_path": dataset_folder, | |
| "caption_ext": "txt", | |
| "caption_dropout_rate": 0.05, | |
| "shuffle_tokens": bool(shuffle_tokens), | |
| "cache_latents_to_disk": bool(cache_latents), | |
| "resolution": resolutions, | |
| } | |
| ], | |
| "train": train_cfg, | |
| "model": model_cfg, | |
| "sample": { | |
| "sampler": sampler, | |
| "sample_every": int(sample_every), | |
| "sample_start_step": 0, | |
| "width": max(resolutions), | |
| "height": max(resolutions), | |
| "prompts": valid_prompts if valid_prompts else ["1girl, solo, masterpiece, best quality"], | |
| "neg": negative_prompt or "", | |
| "seed": 42, | |
| "walk_seed": True, | |
| "guidance_scale": float(guidance_scale), | |
| "sample_steps": int(sample_steps), | |
| }, | |
| } | |
| if trigger_word and trigger_word.strip(): | |
| process_dict["trigger_word"] = trigger_word.strip() | |
| if use_advanced_yaml and advanced_yaml_str and advanced_yaml_str.strip(): | |
| try: | |
| extra_cfg = yaml.safe_load(advanced_yaml_str) | |
| if isinstance(extra_cfg, dict): | |
| process_dict = recursive_update(process_dict, extra_cfg) | |
| except Exception as e: | |
| raise gr.Error(f"Invalid YAML in Advanced Overrides: {e}") | |
| full_config = { | |
| "job": "extension", | |
| "config": { | |
| "name": slugged_name, | |
| "process": [process_dict], | |
| }, | |
| "meta": { | |
| "name": "[name]", | |
| "version": "1.0", | |
| }, | |
| } | |
| os.makedirs("tmp", exist_ok=True) | |
| config_path = os.path.join("tmp", f"job_{slugged_name}.yaml") | |
| with open(config_path, "w", encoding="utf-8") as f: | |
| yaml.dump(full_config, f, sort_keys=False) | |
| return full_config, config_path | |
| # ============================================================================== | |
| # ============================================================================== | |
| # ZeroGPU Time Allocation + CPU-side Estimator (launch-calcium style) | |
| # ============================================================================== | |
| IS_SPACE = bool(os.environ.get("SPACE_ID")) | |
| DEFAULT_GPU_SECONDS = 120 | |
| DEFAULT_INF_GPU_SECONDS = 60 | |
| TIMING_LOG = os.environ.get("TIMING_LOG") or os.path.join( | |
| tempfile.gettempdir(), "zerogpu_trainer_timings.json" | |
| ) | |
| TIMING_LOG_MAX_ROWS = 200 | |
| # Fallback training rate (seconds per step*batch*megapixel) + fixed startup overhead (model load + latent cache) | |
| FALLBACK_TRAIN_RATE = 1.45 | |
| FALLBACK_TRAIN_OVERHEAD = 26.0 | |
| # Fallback inference rate (seconds per step*megapixel) + overhead | |
| FALLBACK_INF_RATE = 0.45 | |
| FALLBACK_INF_OVERHEAD = 6.0 | |
| _timing_lock = threading.Lock() | |
| def _load_timings() -> list: | |
| try: | |
| with open(TIMING_LOG, "r", encoding="utf-8") as fh: | |
| rows = json.load(fh) | |
| return rows if isinstance(rows, list) else [] | |
| except (OSError, ValueError): | |
| return [] | |
| def _log_timing( | |
| mode: str, | |
| mpx: float, | |
| steps: int, | |
| batch: int, | |
| seconds: float, | |
| arch: str = "sdxl", | |
| ) -> None: | |
| row = { | |
| "mode": str(mode), | |
| "arch": str(arch), | |
| "mpx": round(float(mpx), 4), | |
| "steps": int(steps), | |
| "batch": int(batch), | |
| "seconds": round(float(seconds), 3), | |
| } | |
| try: | |
| with _timing_lock: | |
| rows = _load_timings() | |
| rows.append(row) | |
| rows = rows[-TIMING_LOG_MAX_ROWS:] | |
| with open(TIMING_LOG, "w", encoding="utf-8") as fh: | |
| json.dump(rows, fh) | |
| except OSError as exc: | |
| print(f"[timing] could not write log: {exc}") | |
| def estimate_train_gpu_seconds( | |
| chunk_steps, | |
| resolution_choice, | |
| batch_size=1, | |
| grad_accum=1, | |
| preset_name="Illustrious XL v0.1 (Diffusers Repo)", | |
| train_text_encoder=True, | |
| disable_sampling=False, | |
| sample_steps=28, | |
| ): | |
| """Pure-CPU estimate of how long one training chunk will take on ZeroGPU. | |
| Model: seconds ~= overhead + rate * (chunk_steps * batch * grad_accum * megapixels). | |
| Fitted from the rolling log of measured training chunks (with linear recency | |
| weighting), falling back to calibrated Illustrious/SDXL/FLUX constants when | |
| no history exists yet. | |
| """ | |
| steps = max(1, int(chunk_steps or 1)) | |
| batch = max(1, int(batch_size or 1)) * max(1, int(grad_accum or 1)) | |
| res_list = [int(r) for r in (resolution_choice or ["1024"])] if resolution_choice else [1024] | |
| max_res = max(res_list) if res_list else 1024 | |
| mpx = max(0.05, (max_res * max_res) / 1_000_000) | |
| is_flux = "FLUX" in str(preset_name) | |
| arch_key = "flux" if is_flux else "sdxl" | |
| te_mult = 1.25 if train_text_encoder else 1.0 | |
| work = steps * batch * mpx * te_mult | |
| sample_overhead = 0.0 | |
| if not disable_sampling: | |
| sample_overhead = 2.0 * max(1, int(sample_steps or 20)) * mpx * 0.35 | |
| rows = [ | |
| r | |
| for r in _load_timings() | |
| if r.get("mode") == "train" | |
| and r.get("arch", "sdxl") == arch_key | |
| and r.get("seconds", 0) > 0 | |
| and r.get("steps", 0) > 0 | |
| and r.get("mpx", 0) > 0 | |
| ] | |
| base_overhead = FALLBACK_TRAIN_OVERHEAD * (1.35 if is_flux else 1.0) | |
| base_rate = FALLBACK_TRAIN_RATE * (1.6 if is_flux else 1.0) | |
| if not rows: | |
| return int(round(base_overhead + base_rate * work + sample_overhead)) | |
| smallest = min(rows, key=lambda r: max(1, r.get("batch", 1)) * r["steps"] * r["mpx"]) | |
| overhead = min(base_overhead, smallest["seconds"] * 0.5) | |
| weighted_sum = 0.0 | |
| weight_total = 0.0 | |
| for index, row in enumerate(rows, start=1): | |
| row_work = max(1, row.get("batch", 1)) * row["steps"] * row["mpx"] | |
| if row_work <= 0: | |
| continue | |
| rate = max(0.0, (row["seconds"] - overhead)) / row_work | |
| weight = index | |
| weighted_sum += rate * weight | |
| weight_total += weight | |
| rate = (weighted_sum / weight_total) if weight_total else base_rate | |
| if rate <= 0: | |
| rate = base_rate | |
| return int(round(overhead + rate * work + sample_overhead)) | |
| def gpu_train_time_check( | |
| chunk_steps, | |
| resolution_choice, | |
| batch_size, | |
| grad_accum, | |
| preset_name, | |
| train_text_encoder, | |
| disable_sampling, | |
| sample_steps, | |
| ): | |
| """Handler for the 'Gpu time check' button on the Training tab — never touches the GPU.""" | |
| return estimate_train_gpu_seconds( | |
| chunk_steps=chunk_steps, | |
| resolution_choice=resolution_choice, | |
| batch_size=batch_size, | |
| grad_accum=grad_accum, | |
| preset_name=preset_name, | |
| train_text_encoder=train_text_encoder, | |
| disable_sampling=disable_sampling, | |
| sample_steps=sample_steps, | |
| ) | |
| def estimate_inf_gpu_seconds(steps, guidance_scale, width, height): | |
| """Pure-CPU estimate of how long inference settings will take on ZeroGPU.""" | |
| steps = max(1, int(steps or 1)) | |
| mpx = max(0.01, (int(width or 1024) * int(height or 1024)) / 1_000_000) | |
| work = steps * mpx | |
| rows = [ | |
| r | |
| for r in _load_timings() | |
| if r.get("mode") == "inference" and r.get("seconds", 0) > 0 and r.get("steps", 0) > 0 and r.get("mpx", 0) > 0 | |
| ] | |
| if not rows: | |
| return int(round(FALLBACK_INF_OVERHEAD + FALLBACK_INF_RATE * work)) | |
| smallest = min(rows, key=lambda r: r["steps"] * r["mpx"]) | |
| overhead = min(FALLBACK_INF_OVERHEAD, smallest["seconds"] * 0.25) | |
| weighted_sum = 0.0 | |
| weight_total = 0.0 | |
| for index, row in enumerate(rows, start=1): | |
| row_work = row["steps"] * row["mpx"] | |
| if row_work <= 0: | |
| continue | |
| rate = max(0.0, (row["seconds"] - overhead)) / row_work | |
| weight = index | |
| weighted_sum += rate * weight | |
| weight_total += weight | |
| rate = (weighted_sum / weight_total) if weight_total else FALLBACK_INF_RATE | |
| if rate <= 0: | |
| rate = FALLBACK_INF_RATE | |
| return int(round(overhead + rate * work)) | |
| def gpu_inf_time_check(steps, guidance_scale, width, height): | |
| """Handler for the 'Gpu time check' button on the Inference tab — never touches the GPU.""" | |
| return estimate_inf_gpu_seconds(steps, guidance_scale, width, height) | |
| def _get_train_duration( | |
| config_path: str, | |
| hf_token: Optional[str] = None, | |
| gpu_seconds: int = DEFAULT_GPU_SECONDS, | |
| ) -> int: | |
| """Dynamic ZeroGPU allocation for a training chunk: whatever the user typed in the allocate box.""" | |
| try: | |
| requested = int(gpu_seconds) | |
| except (TypeError, ValueError): | |
| return DEFAULT_GPU_SECONDS | |
| return max(1, requested) | |
| def _get_smoke_duration( | |
| preset_name: str = "Illustrious XL v0.1 (Diffusers Repo)", | |
| custom_model_path: str = "", | |
| custom_arch: str = "sdxl", | |
| steps: int = 5, | |
| hf_token: Optional[str] = None, | |
| gpu_seconds: int = DEFAULT_GPU_SECONDS, | |
| ) -> int: | |
| """Dynamic ZeroGPU allocation for the smoke test.""" | |
| try: | |
| requested = int(gpu_seconds) | |
| except (TypeError, ValueError): | |
| return DEFAULT_GPU_SECONDS | |
| return max(1, requested) | |
| def _get_inf_duration( | |
| base_model_preset: str, | |
| custom_base_model: str, | |
| lora_weights_path: str, | |
| prompt: str, | |
| negative_prompt: str, | |
| width: int, | |
| height: int, | |
| steps: int, | |
| guidance_scale: float, | |
| lora_scale: float, | |
| seed: int, | |
| hf_token: Optional[str] = None, | |
| gpu_seconds: int = DEFAULT_INF_GPU_SECONDS, | |
| ) -> int: | |
| """Dynamic ZeroGPU allocation for inference.""" | |
| try: | |
| requested = int(gpu_seconds) | |
| except (TypeError, ValueError): | |
| return DEFAULT_INF_GPU_SECONDS | |
| return max(1, requested) | |
| # ZeroGPU (`@spaces.GPU`) Decorated Functions | |
| # ============================================================================== | |
| def gpu_caption_images( | |
| image_paths: List[str], | |
| trigger_word: str, | |
| caption_style: str, | |
| existing_captions: List[str], | |
| overwrite_existing: bool, | |
| ) -> List[str]: | |
| """ | |
| Runs Florence-2 on ZeroGPU to caption uploaded images. | |
| Supports natural language descriptions (for FLUX/SDXL) and comma-separated tags (for Illustrious XL). | |
| """ | |
| from transformers import AutoProcessor, AutoModelForCausalLM | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| torch_dtype = torch.float16 if device == "cuda" else torch.float32 | |
| model_id = "multimodalart/Florence-2-large-no-flash-attn" | |
| model = AutoModelForCausalLM.from_pretrained( | |
| model_id, torch_dtype=torch_dtype, trust_remote_code=True | |
| ).to(device) | |
| processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True) | |
| task_map = { | |
| "Detailed Natural Language (<MORE_DETAILED_CAPTION>)": "<MORE_DETAILED_CAPTION>", | |
| "Standard Caption (<DETAILED_CAPTION>)": "<DETAILED_CAPTION>", | |
| "Concise / Tag-like (<CAPTION>)": "<CAPTION>", | |
| "Danbooru / Illustrious Tag Style (Formatted)": "<MORE_DETAILED_CAPTION>", | |
| } | |
| task_prompt = task_map.get(caption_style, "<DETAILED_CAPTION>") | |
| is_tag_style = "Tag Style" in caption_style | |
| updated_captions = list(existing_captions) | |
| for i, img_path in enumerate(image_paths): | |
| if not img_path or not os.path.exists(img_path): | |
| continue | |
| if not overwrite_existing and i < len(updated_captions) and updated_captions[i] and updated_captions[i].strip() and updated_captions[i].strip() != "[trigger]": | |
| continue | |
| image = Image.open(img_path).convert("RGB") | |
| inputs = processor(text=task_prompt, images=image, return_tensors="pt").to(device, torch_dtype) | |
| with torch.no_grad(): | |
| generated_ids = model.generate( | |
| input_ids=inputs["input_ids"], | |
| pixel_values=inputs["pixel_values"], | |
| max_new_tokens=512, | |
| num_beams=3, | |
| ) | |
| generated_text = processor.batch_decode(generated_ids, skip_special_tokens=False)[0] | |
| parsed = processor.post_process_generation( | |
| generated_text, task=task_prompt, image_size=(image.width, image.height) | |
| ) | |
| raw_caption = parsed.get(task_prompt, "").strip() | |
| for prefix in ("The image shows ", "This image shows ", "The image depicts ", "An illustration of ", "A digital illustration of "): | |
| if raw_caption.startswith(prefix): | |
| raw_caption = raw_caption[len(prefix):] | |
| if is_tag_style: | |
| # Convert clauses into clean comma-separated tags suitable for Illustrious XL | |
| parts = re.split(r"[.,;]+|\band\b|\bwith\b", raw_caption) | |
| tags = [] | |
| for p in parts: | |
| cleaned = re.sub(r"^(a|an|the|her|his|their|its|she|he|is|are|has|wearing|in)\s+", "", p.strip(), flags=re.IGNORECASE) | |
| cleaned = cleaned.strip() | |
| if len(cleaned) > 2 and cleaned.lower() not in tags: | |
| tags.append(cleaned) | |
| raw_caption = "masterpiece, best quality, newest, " + ", ".join(tags) | |
| if trigger_word and trigger_word.strip(): | |
| trig = trigger_word.strip() | |
| if trig not in raw_caption and "[trigger]" not in raw_caption: | |
| raw_caption = f"{trig}, {raw_caption}" if is_tag_style else f"{trig} {raw_caption}" | |
| if i < len(updated_captions): | |
| updated_captions[i] = raw_caption | |
| else: | |
| updated_captions.append(raw_caption) | |
| model.to("cpu") | |
| del model | |
| del processor | |
| gc.collect() | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| return updated_captions | |
| def gpu_train_chunk(config_path: str, hf_token: Optional[str] = None, gpu_seconds: int = DEFAULT_GPU_SECONDS) -> Dict[str, Any]: | |
| """ | |
| Executes a single training chunk on ZeroGPU (`@spaces.GPU`). | |
| Automatically loads the latest checkpoint + optimizer.pt + ema.pt from `save_root`, | |
| trains up to `max_steps_per_run` (or `steps`), saves the new checkpoint, and frees VRAM. | |
| """ | |
| started = time.time() | |
| if hf_token: | |
| os.environ["HF_TOKEN"] = hf_token | |
| # Reset global accelerator inside the forked ZeroGPU worker so it binds to CUDA | |
| reset_accelerator() | |
| with open(config_path, "r", encoding="utf-8") as f: | |
| cfg = yaml.safe_load(f) | |
| lora_name = cfg["config"]["name"] | |
| proc_cfg = cfg["config"]["process"][0] | |
| training_folder = proc_cfg.get("training_folder", "output") | |
| total_steps = int(proc_cfg["train"]["steps"]) | |
| prev_step, _, _ = get_checkpoint_status(training_folder, lora_name) | |
| job = get_job(config_path) | |
| try: | |
| job.run() | |
| finally: | |
| job.cleanup() | |
| del job | |
| gc.collect() | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| elapsed = time.time() - started | |
| current_step, latest_ckpt, sample_imgs = get_checkpoint_status(training_folder, lora_name) | |
| steps_ran = max(1, current_step - prev_step) | |
| res_list = proc_cfg.get("datasets", [{}])[0].get("resolution", [1024]) | |
| max_res = max(res_list) if isinstance(res_list, list) and res_list else int(res_list or 1024) | |
| batch_eff = int(proc_cfg["train"].get("batch_size", 1)) * int(proc_cfg["train"].get("gradient_accumulation_steps", 1)) | |
| arch_str = "flux" if proc_cfg.get("model", {}).get("is_flux") else "sdxl" | |
| _log_timing( | |
| mode="train", | |
| mpx=(max_res * max_res) / 1_000_000, | |
| steps=steps_ran, | |
| batch=batch_eff, | |
| seconds=elapsed, | |
| arch=arch_str, | |
| ) | |
| return { | |
| "current_step": current_step, | |
| "total_steps": total_steps, | |
| "latest_ckpt": latest_ckpt, | |
| "sample_imgs": sample_imgs, | |
| "is_complete": current_step >= total_steps, | |
| "elapsed_seconds": round(elapsed, 1), | |
| "allocated_seconds": int(gpu_seconds), | |
| } | |
| def gpu_quick_smoke_test( | |
| preset_name: str = "Illustrious XL v0.1 (Diffusers Repo)", | |
| custom_model_path: str = "", | |
| custom_arch: str = "sdxl", | |
| steps: int = 5, | |
| hf_token: Optional[str] = None, | |
| gpu_seconds: int = DEFAULT_GPU_SECONDS, | |
| ) -> Tuple[str, Optional[str]]: | |
| """ | |
| Runs a fast end-to-end ZeroGPU smoke test (maximum 4 minutes / 240s on ZeroGPU): | |
| generates a synthetic 2-image dataset, runs `steps` training steps via ai-toolkit, | |
| and verifies the output `.safetensors` LoRA checkpoint. | |
| """ | |
| t0 = time.time() | |
| token = resolve_hf_token(hf_token) | |
| reset_accelerator() | |
| test_id = f"smoke_test_{int(t0)}" | |
| test_ds_dir = os.path.join("tmp", test_id, "dataset") | |
| test_out_dir = os.path.join("tmp", test_id, "output") | |
| os.makedirs(test_ds_dir, exist_ok=True) | |
| os.makedirs(test_out_dir, exist_ok=True) | |
| # Create 2 synthetic training images + .txt captions | |
| colors = [((60, 90, 180), "1girl, solo, blue theme, masterpiece, best quality"), | |
| ((180, 70, 90), "1girl, solo, red theme, masterpiece, best quality")] | |
| img_paths = [] | |
| caps = [] | |
| for idx, (col, cap) in enumerate(colors): | |
| p = os.path.join(test_ds_dir, f"synth_{idx}.png") | |
| img = Image.new("RGB", (512, 512), color=col) | |
| img.save(p) | |
| img_paths.append(p) | |
| caps.append(cap) | |
| final_ds = create_dataset_folder(img_paths, caps, dataset_dir=os.path.join(test_ds_dir, "prepared")) | |
| _, config_path = build_training_config( | |
| lora_name=test_id, | |
| trigger_word="smoketest_tok", | |
| preset_name=preset_name, | |
| custom_model_path=custom_model_path, | |
| custom_arch=custom_arch, | |
| dataset_folder=final_ds, | |
| total_steps=int(steps), | |
| chunk_steps=int(steps), | |
| lr=1e-4, | |
| unet_lr=1e-4, | |
| text_encoder_lr=5e-5, | |
| train_text_encoder=False, | |
| rank=4, | |
| alpha=4, | |
| batch_size=1, | |
| grad_accum=1, | |
| resolution_choice=["512"], | |
| optimizer_type="adamw8bit", | |
| lr_scheduler_type="constant", | |
| quantize_model="FLUX" in preset_name, | |
| quantize_te="FLUX" in preset_name, | |
| low_vram=False, | |
| use_ema=False, | |
| cache_latents=True, | |
| cache_text_embeddings=False, | |
| shuffle_tokens=True, | |
| save_every=int(steps), | |
| sample_every=int(steps) + 100, | |
| disable_sampling=True, | |
| sample_prompts=[], | |
| negative_prompt="", | |
| guidance_scale=6.0, | |
| sample_steps=10, | |
| push_to_hub=False, | |
| push_every_chunk=False, | |
| hf_repo_id="", | |
| hf_private=True, | |
| use_advanced_yaml=False, | |
| advanced_yaml_str="", | |
| training_folder=test_out_dir, | |
| ) | |
| job = get_job(config_path) | |
| try: | |
| job.run() | |
| finally: | |
| job.cleanup() | |
| del job | |
| gc.collect() | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| elapsed = time.time() - t0 | |
| _log_timing( | |
| mode="train", | |
| mpx=(512 * 512) / 1_000_000, | |
| steps=int(steps), | |
| batch=1, | |
| seconds=elapsed, | |
| arch="flux" if "FLUX" in preset_name else "sdxl", | |
| ) | |
| step_reached, ckpt_path, _ = get_checkpoint_status(test_out_dir, test_id) | |
| device_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU" | |
| summary = ( | |
| f"✅ **ZeroGPU Test Run Succeeded in {elapsed:.1f}s (Allocated: {int(gpu_seconds)}s)**\n" | |
| f"- **Device**: `{device_name}`\n" | |
| f"- **Model Preset**: `{preset_name}`\n" | |
| f"- **Steps Completed**: `{step_reached}/{steps}`\n" | |
| f"- **Saved Checkpoint**: `{ckpt_path}`" | |
| ) | |
| return summary, ckpt_path | |
| def gpu_test_inference( | |
| base_model_preset: str, | |
| custom_base_model: str, | |
| lora_weights_path: str, | |
| prompt: str, | |
| negative_prompt: str, | |
| width: int, | |
| height: int, | |
| steps: int, | |
| guidance_scale: float, | |
| lora_scale: float, | |
| seed: int, | |
| hf_token: Optional[str] = None, | |
| gpu_seconds: int = DEFAULT_INF_GPU_SECONDS, | |
| ) -> Image.Image: | |
| """ | |
| Runs quick inference on ZeroGPU to test a trained Illustrious XL / SDXL / FLUX LoRA. | |
| """ | |
| started = time.time() | |
| if hf_token: | |
| os.environ["HF_TOKEN"] = hf_token | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| dtype = torch.bfloat16 if (device == "cuda" and torch.cuda.is_bf16_supported()) else (torch.float16 if device == "cuda" else torch.float32) | |
| preset = MODEL_PRESETS.get(base_model_preset, MODEL_PRESETS["Illustrious XL v0.1 (Diffusers Repo)"]) | |
| model_id = custom_base_model.strip() if (custom_base_model and custom_base_model.strip()) else preset["name_or_path"] | |
| is_flux = preset.get("is_flux", False) or "flux" in model_id.lower() | |
| if is_flux: | |
| from diffusers import FluxPipeline | |
| pipe = FluxPipeline.from_pretrained(model_id, torch_dtype=dtype) | |
| else: | |
| from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler | |
| if model_id.lower().endswith(".safetensors"): | |
| if not os.path.exists(model_id) and model_id.count("/") >= 2: | |
| parts = model_id.split("/") | |
| model_id = hf_hub_download(repo_id="/".join(parts[:2]), filename="/".join(parts[2:])) | |
| pipe = StableDiffusionXLPipeline.from_single_file(model_id, torch_dtype=dtype) | |
| else: | |
| pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=dtype) | |
| pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config) | |
| pipe.to(device) | |
| if lora_weights_path and str(lora_weights_path).strip(): | |
| lpath = str(lora_weights_path).strip() | |
| if os.path.exists(lpath): | |
| pipe.load_lora_weights(os.path.dirname(lpath) or ".", weight_name=os.path.basename(lpath)) | |
| else: | |
| pipe.load_lora_weights(lpath) | |
| gen = torch.Generator(device=device).manual_seed(int(seed) if seed >= 0 else random.randint(0, 2**31 - 1)) | |
| kwargs: Dict[str, Any] = { | |
| "prompt": prompt, | |
| "width": int(width), | |
| "height": int(height), | |
| "num_inference_steps": int(steps), | |
| "guidance_scale": float(guidance_scale), | |
| "generator": gen, | |
| } | |
| if not is_flux and negative_prompt: | |
| kwargs["negative_prompt"] = negative_prompt | |
| if lora_weights_path and str(lora_weights_path).strip(): | |
| kwargs["joint_attention_kwargs" if is_flux else "cross_attention_kwargs"] = {"scale": float(lora_scale)} | |
| with torch.no_grad(): | |
| image = pipe(**kwargs).images[0] | |
| pipe.to("cpu") | |
| del pipe | |
| gc.collect() | |
| if torch.cuda.is_available(): | |
| torch.cuda.empty_cache() | |
| _log_timing( | |
| mode="inference", | |
| mpx=(int(width) * int(height)) / 1_000_000, | |
| steps=int(steps), | |
| batch=1, | |
| seconds=time.time() - started, | |
| arch="flux" if is_flux else "sdxl", | |
| ) | |
| return image | |
| # ============================================================================== | |
| # Gradio UI Event Handlers | |
| # ============================================================================== | |
| def on_preset_change(preset_name: str): | |
| preset = MODEL_PRESETS.get(preset_name, MODEL_PRESETS["Illustrious XL v0.1 (Diffusers Repo)"]) | |
| is_custom = preset_name == "Custom Diffusers / HF Hub Model" | |
| return ( | |
| gr.update( | |
| value=preset["name_or_path"], | |
| visible=True, | |
| label="Model Repo ID or Single-File Path on Hugging Face Hub" if is_custom else "Base Model Path (Override Optional)", | |
| ), | |
| gr.update(visible=is_custom, value=preset["arch"]), | |
| gr.update(value=preset["lr"]), | |
| gr.update(value=preset["unet_lr"]), | |
| gr.update(value=preset["text_encoder_lr"]), | |
| gr.update(value=preset["train_text_encoder"]), | |
| gr.update(value=preset["quantize"]), | |
| gr.update(value=preset["quantize_te"]), | |
| gr.update(value=[str(r) for r in preset["resolution"]]), | |
| gr.update(value=preset["shuffle_tokens"]), | |
| gr.update(value=preset["guidance_scale"]), | |
| gr.update(value=preset["sample_steps"]), | |
| gr.update(value=preset["default_neg"]), | |
| f"ℹ️ **{preset_name}**: {preset['description']}", | |
| ) | |
| def handle_file_upload(uploaded_files, trigger_word): | |
| image_paths, captions_map = extract_uploaded_files(uploaded_files) | |
| if len(image_paths) < 2: | |
| raise gr.Error("Please upload at least 2 images (or a .zip archive containing images) to train your LoRA.") | |
| if len(image_paths) > MAX_IMAGES: | |
| raise gr.Error(f"Maximum {MAX_IMAGES} images allowed per dataset.") | |
| updates = [image_paths, gr.update(visible=True)] | |
| for i in range(1, MAX_IMAGES + 1): | |
| visible = i <= len(image_paths) | |
| img_val = image_paths[i - 1] if visible else None | |
| cap_val = "" | |
| if visible and img_val: | |
| stem = os.path.splitext(os.path.basename(img_val))[0] | |
| if stem in captions_map: | |
| cap_val = captions_map[stem] | |
| elif trigger_word and trigger_word.strip(): | |
| cap_val = f"{trigger_word.strip()}" | |
| updates.append(gr.update(visible=visible)) | |
| updates.append(gr.update(value=img_val, visible=visible)) | |
| updates.append(gr.update(value=cap_val, visible=visible)) | |
| trig_display = trigger_word.strip() if (trigger_word and trigger_word.strip()) else "1girl, solo" | |
| updates.append(gr.update(visible=True)) | |
| updates.append(gr.update(value=f"{trig_display}, masterpiece, best quality, newest, detailed background")) | |
| updates.append(gr.update(value=f"{trig_display}, portrait, dramatic lighting, looking at viewer")) | |
| updates.append(gr.update(value="")) | |
| updates.append(gr.update(visible=True)) | |
| return updates | |
| def handle_captioning_click(image_paths, trigger_word, caption_style, overwrite_existing, *current_captions): | |
| if not image_paths: | |
| raise gr.Error("Please upload images first.") | |
| active_captions = list(current_captions[:len(image_paths)]) | |
| new_captions = gpu_caption_images( | |
| image_paths=image_paths, | |
| trigger_word=trigger_word or "", | |
| caption_style=caption_style, | |
| existing_captions=active_captions, | |
| overwrite_existing=bool(overwrite_existing), | |
| ) | |
| padded = new_captions + list(current_captions[len(new_captions):]) | |
| return padded[:MAX_IMAGES] | |
| def run_chunked_training_pipeline( | |
| image_paths: List[str], | |
| lora_name: str, | |
| trigger_word: str, | |
| preset_name: str, | |
| custom_model_path: str, | |
| custom_arch: str, | |
| total_steps: int, | |
| chunk_steps: int, | |
| auto_loop_chunks: bool, | |
| lr: float, | |
| unet_lr: float, | |
| text_encoder_lr: float, | |
| train_text_encoder: bool, | |
| rank: int, | |
| alpha: int, | |
| batch_size: int, | |
| grad_accum: int, | |
| resolution_choice: List[str], | |
| optimizer_type: str, | |
| lr_scheduler_type: str, | |
| quantize_model: bool, | |
| quantize_te: bool, | |
| low_vram: bool, | |
| use_ema: bool, | |
| cache_latents: bool, | |
| cache_text_embeddings: bool, | |
| shuffle_tokens: bool, | |
| save_every: int, | |
| sample_every: int, | |
| disable_sampling: bool, | |
| sample_1: str, | |
| sample_2: str, | |
| sample_3: str, | |
| negative_prompt: str, | |
| guidance_scale: float, | |
| sample_steps: int, | |
| push_to_hub: bool, | |
| push_every_chunk: bool, | |
| hf_repo_id_input: str, | |
| hf_private: bool, | |
| resume_from_hub: bool, | |
| hf_token_input: str, | |
| use_advanced_yaml: bool, | |
| advanced_yaml_str: str, | |
| gpu_seconds: int = DEFAULT_GPU_SECONDS, | |
| *captions, | |
| ): | |
| """ | |
| Generator orchestrating chunked ZeroGPU training across multiple `@spaces.GPU` slices. | |
| Runs on CPU between GPU chunks so ZeroGPU timeouts are never exceeded and checkpoints/Hub uploads | |
| happen reliably after every chunk. | |
| """ | |
| if not lora_name or not lora_name.strip(): | |
| raise gr.Error("Please provide a unique name for your LoRA.") | |
| if not image_paths or len(image_paths) < 2: | |
| raise gr.Error("Please upload at least 2 training images before starting training.") | |
| slugged_name = slugify(lora_name.strip()) | |
| training_folder = "output" | |
| token = resolve_hf_token(hf_token_input) | |
| # Resolve HF Hub repository ID if pushing or resuming from Hub | |
| hf_repo_id = hf_repo_id_input.strip() if (hf_repo_id_input and hf_repo_id_input.strip()) else "" | |
| if (push_to_hub or resume_from_hub) and not hf_repo_id: | |
| username = get_hf_username(token) | |
| if username: | |
| hf_repo_id = f"{username}/{slugged_name}" | |
| elif push_to_hub: | |
| push_to_hub = False | |
| gr.Warning("No valid Hugging Face write token detected. Training will save checkpoints locally in the Space.") | |
| log_lines: List[str] = [] | |
| def _format_log(msg: str) -> str: | |
| ts = time.strftime("%H:%M:%S") | |
| log_lines.append(f"[{ts}] {msg}") | |
| return "\n\n".join(log_lines) | |
| # 1. Build dataset folder with .txt captions and metadata.jsonl | |
| dataset_dir = os.path.join("datasets", slugged_name) | |
| active_captions = list(captions[:len(image_paths)]) | |
| create_dataset_folder(image_paths, active_captions, dataset_dir=dataset_dir) | |
| yield _format_log(f"✅ Prepared dataset with **{len(image_paths)}** images and `.txt` captions in `{dataset_dir}`."), None, [] | |
| # 2. Optionally pull existing checkpoint from Hugging Face Hub for resumption | |
| if resume_from_hub and hf_repo_id: | |
| try: | |
| yield _format_log(f"🔄 Checking Hugging Face Hub `{hf_repo_id}` for existing checkpoints to resume..."), None, [] | |
| hub_msg = sync_checkpoint_from_hub(hf_repo_id, training_folder, slugged_name, token=token) | |
| yield _format_log(f"📥 {hub_msg}"), None, [] | |
| except Exception as e: | |
| yield _format_log(f"ℹ️ Could not pull existing checkpoint from `{hf_repo_id}` ({e}). Starting fresh or from local storage."), None, [] | |
| # 3. Check current step status | |
| current_step, latest_ckpt, sample_imgs = get_checkpoint_status(training_folder, slugged_name) | |
| total_steps = int(total_steps) | |
| chunk_steps = int(chunk_steps) if (chunk_steps and int(chunk_steps) > 0) else total_steps | |
| if current_step >= total_steps: | |
| yield ( | |
| _format_log(f"🎉 Training is already at step **{current_step}/{total_steps}**! Increase Total Steps to train further, or download your checkpoint below."), | |
| latest_ckpt, | |
| sample_imgs, | |
| ) | |
| return | |
| # 4. Generate YAML config | |
| _, config_path = build_training_config( | |
| lora_name=slugged_name, | |
| trigger_word=trigger_word, | |
| preset_name=preset_name, | |
| custom_model_path=custom_model_path, | |
| custom_arch=custom_arch, | |
| dataset_folder=dataset_dir, | |
| total_steps=total_steps, | |
| chunk_steps=chunk_steps, | |
| lr=lr, | |
| unet_lr=unet_lr, | |
| text_encoder_lr=text_encoder_lr, | |
| train_text_encoder=train_text_encoder, | |
| rank=rank, | |
| alpha=alpha, | |
| batch_size=batch_size, | |
| grad_accum=grad_accum, | |
| resolution_choice=resolution_choice, | |
| optimizer_type=optimizer_type, | |
| lr_scheduler_type=lr_scheduler_type, | |
| quantize_model=quantize_model, | |
| quantize_te=quantize_te, | |
| low_vram=low_vram, | |
| use_ema=use_ema, | |
| cache_latents=cache_latents, | |
| cache_text_embeddings=cache_text_embeddings, | |
| shuffle_tokens=shuffle_tokens, | |
| save_every=save_every, | |
| sample_every=sample_every, | |
| disable_sampling=disable_sampling, | |
| sample_prompts=[sample_1, sample_2, sample_3], | |
| negative_prompt=negative_prompt, | |
| guidance_scale=guidance_scale, | |
| sample_steps=sample_steps, | |
| push_to_hub=push_to_hub, | |
| push_every_chunk=push_every_chunk, | |
| hf_repo_id=hf_repo_id, | |
| hf_private=hf_private, | |
| use_advanced_yaml=use_advanced_yaml, | |
| advanced_yaml_str=advanced_yaml_str, | |
| training_folder=training_folder, | |
| ) | |
| # 5. Iterative Chunked ZeroGPU Training Loop | |
| chunk_idx = 0 | |
| while current_step < total_steps: | |
| chunk_idx += 1 | |
| next_target = min(total_steps, current_step + chunk_steps) | |
| alloc_sec = max(1, int(gpu_seconds or DEFAULT_GPU_SECONDS)) | |
| yield ( | |
| _format_log( | |
| f"🚀 **ZeroGPU Chunk #{chunk_idx}** (Allocated `{alloc_sec}s` GPU time): Training steps **{current_step} → {next_target}** (of {total_steps} total)..." | |
| ), | |
| latest_ckpt, | |
| sample_imgs, | |
| ) | |
| result = gpu_train_chunk(config_path=config_path, hf_token=token, gpu_seconds=alloc_sec) | |
| current_step = result["current_step"] | |
| latest_ckpt = result["latest_ckpt"] | |
| sample_imgs = result["sample_imgs"] | |
| elapsed_s = result.get("elapsed_seconds", 0.0) | |
| yield ( | |
| _format_log( | |
| f"💾 **Chunk #{chunk_idx} Complete in {elapsed_s}s / {alloc_sec}s**: Reached step **{current_step}/{total_steps}**. " | |
| f"Latest checkpoint: `{os.path.basename(latest_ckpt) if latest_ckpt else 'N/A'}`" | |
| ), | |
| latest_ckpt, | |
| sample_imgs, | |
| ) | |
| if not auto_loop_chunks or current_step >= total_steps: | |
| break | |
| if current_step >= total_steps: | |
| hub_link = f" and uploaded to [https://huggingface.co/{hf_repo_id}](https://huggingface.co/{hf_repo_id})" if (push_to_hub and hf_repo_id) else "" | |
| yield ( | |
| _format_log(f"🎉 **Training Completed ({current_step}/{total_steps} steps)!** Model saved as `{slugged_name}.safetensors`{hub_link}."), | |
| latest_ckpt, | |
| sample_imgs, | |
| ) | |
| else: | |
| yield ( | |
| _format_log(f"⏸️ **Paused after 1 chunk at step {current_step}/{total_steps}**. Click **Start / Resume Chunked Training** to continue from step {current_step}."), | |
| latest_ckpt, | |
| sample_imgs, | |
| ) | |
| # ============================================================================== | |
| # Build Gradio UI | |
| # ============================================================================== | |
| DEFAULT_ADVANCED_YAML = """# Optional YAML overrides merged into config.process[0] | |
| # Example: | |
| # train: | |
| # noise_offset: 0.03 | |
| """ | |
| theme = gr.themes.Soft( | |
| primary_hue="indigo", | |
| secondary_hue="slate", | |
| font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"], | |
| ) | |
| css = """ | |
| .container { max-width: 1280px; margin: auto; } | |
| .header-box { padding: 12px 16px; border-radius: 10px; margin-bottom: 12px; } | |
| .unit-note { display: flex; align-items: center; opacity: 0.75; } | |
| """ | |
| def build_ui() -> gr.Blocks: | |
| with gr.Blocks(theme=theme, css=css, title="Ostris AI-Toolkit on ZeroGPU") as demo: | |
| gr.Markdown( | |
| """ | |
| # ⚡ Ostris AI-Toolkit on ZeroGPU | |
| ### Train **Illustrious XL v0.1**, **FLUX.1 (`dev` / `schnell`)**, **SDXL**, and **Custom Diffusers Models** with Chunked ZeroGPU Checkpointing & Hub Sync | |
| """ | |
| ) | |
| uploaded_image_paths_state = gr.State([]) | |
| with gr.Tabs(): | |
| # ================================================================== | |
| # TAB 1: TRAIN LORA | |
| # ================================================================== | |
| with gr.Tab("🏋️ Train LoRA (ZeroGPU Chunked)"): | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| lora_name = gr.Textbox( | |
| label="LoRA Output Name", | |
| placeholder="e.g. illustrious-style-v1 or my-character-lora", | |
| value="my_illustrious_lora", | |
| info="Used for checkpoint filenames and your Hugging Face Hub repo name.", | |
| ) | |
| with gr.Column(scale=2): | |
| trigger_word = gr.Textbox( | |
| label="Trigger Word / Tag (Optional)", | |
| placeholder="e.g. chr_name, 1girl, or 'in the style of sks'", | |
| value="", | |
| info="Automatically injected into captions or replaces [trigger].", | |
| ) | |
| with gr.Column(scale=2): | |
| hf_token_input = gr.Textbox( | |
| label="Hugging Face Write Token (Optional if HF_TOKEN secret set)", | |
| type="password", | |
| placeholder="hf_...", | |
| info="Needed for pushing checkpoints to HF Hub or downloading gated models like FLUX.1-dev.", | |
| ) | |
| # Model Selection Section | |
| with gr.Group(): | |
| with gr.Row(): | |
| preset_selector = gr.Dropdown( | |
| choices=list(MODEL_PRESETS.keys()), | |
| value="Illustrious XL v0.1 (Diffusers Repo)", | |
| label="Model Architecture / Preset", | |
| scale=2, | |
| ) | |
| custom_model_path = gr.Textbox( | |
| label="Base Model Path (Override Optional)", | |
| value=MODEL_PRESETS["Illustrious XL v0.1 (Diffusers Repo)"]["name_or_path"], | |
| placeholder="org/model-name or org/model-name/file.safetensors", | |
| scale=3, | |
| ) | |
| custom_arch = gr.Dropdown( | |
| choices=CUSTOM_ARCHITECTURES, | |
| value="sdxl", | |
| label="Custom Architecture", | |
| visible=False, | |
| scale=2, | |
| ) | |
| preset_info_md = gr.Markdown( | |
| f"ℹ️ **Illustrious XL v0.1 (Diffusers Repo)**: {MODEL_PRESETS['Illustrious XL v0.1 (Diffusers Repo)']['description']}" | |
| ) | |
| # Dataset Upload & Captioning Section | |
| with gr.Group(): | |
| with gr.Row(): | |
| files_upload = gr.File( | |
| file_types=["image", ".txt", ".zip"], | |
| label="Upload Training Images (.png, .jpg, .webp), .txt Captions, or a .zip Archive", | |
| file_count="multiple", | |
| scale=1, | |
| ) | |
| with gr.Column(scale=3, visible=False) as captioning_area: | |
| gr.Markdown( | |
| "### 🏷️ Dataset Captions & ZeroGPU Auto-Captioning (Florence-2)\n" | |
| "Edit captions individually or auto-generate Danbooru-style tags (for Illustrious XL) / natural language descriptions (for FLUX)." | |
| ) | |
| with gr.Row(): | |
| caption_style = gr.Dropdown( | |
| choices=[ | |
| "Danbooru / Illustrious Tag Style (Formatted)", | |
| "Detailed Natural Language (<MORE_DETAILED_CAPTION>)", | |
| "Standard Caption (<DETAILED_CAPTION>)", | |
| "Concise / Tag-like (<CAPTION>)", | |
| ], | |
| value="Danbooru / Illustrious Tag Style (Formatted)", | |
| label="Florence-2 Captioning Style", | |
| scale=2, | |
| ) | |
| overwrite_captions = gr.Checkbox( | |
| label="Overwrite existing captions", | |
| value=True, | |
| scale=1, | |
| ) | |
| do_captioning_btn = gr.Button( | |
| "✨ Auto-Caption with Florence-2 (ZeroGPU)", | |
| variant="secondary", | |
| scale=2, | |
| ) | |
| output_components = [uploaded_image_paths_state, captioning_area] | |
| caption_boxes = [] | |
| for i in range(1, MAX_IMAGES + 1): | |
| row = gr.Row(visible=False) | |
| with row: | |
| img_comp = gr.Image( | |
| type="filepath", | |
| width=110, | |
| height=110, | |
| min_width=110, | |
| interactive=False, | |
| scale=1, | |
| show_label=False, | |
| ) | |
| txt_comp = gr.Textbox( | |
| label=f"Image #{i} Caption (.txt)", | |
| scale=8, | |
| interactive=True, | |
| lines=2, | |
| ) | |
| output_components.extend([row, img_comp, txt_comp]) | |
| caption_boxes.append(txt_comp) | |
| # ZeroGPU Chunking & Hub Sync Configuration | |
| with gr.Accordion("⚡ ZeroGPU Chunked Training & Hugging Face Hub Sync", open=True): | |
| gr.Markdown( | |
| "ZeroGPU enforces a per-call time limit on `@spaces.GPU`. By splitting training into **chunks** " | |
| "(e.g. 200 steps per GPU call), `ai-toolkit` saves the LoRA `.safetensors`, `optimizer.pt`, and `ema.pt` " | |
| "at the end of each slice and seamlessly resumes on the next slice." | |
| ) | |
| with gr.Row(): | |
| total_steps = gr.Number( | |
| label="Total Training Steps", | |
| value=1000, | |
| minimum=10, | |
| maximum=20000, | |
| step=50, | |
| ) | |
| chunk_steps = gr.Number( | |
| label="Steps per ZeroGPU Chunk", | |
| value=200, | |
| minimum=10, | |
| maximum=2000, | |
| step=25, | |
| info="Number of steps run per @spaces.GPU invocation before checkpointing.", | |
| ) | |
| save_every = gr.Number( | |
| label="Save Checkpoint Every N Steps", | |
| value=200, | |
| minimum=10, | |
| maximum=5000, | |
| step=25, | |
| ) | |
| auto_loop_chunks = gr.Checkbox( | |
| label="Auto-Loop All Chunks Until Total Steps Reached", | |
| value=True, | |
| info="When enabled, automatically chains @spaces.GPU chunks back-to-back.", | |
| ) | |
| # ZeroGPU Time Check & Allocator (launch-calcium style) | |
| with gr.Group(visible=True) as train_gpu_group: | |
| gr.Markdown("**ZeroGPU Time Check & Allocator (per chunk)**") | |
| with gr.Row(): | |
| train_gpu_seconds = gr.Number( | |
| label="Gpu time to allocate", | |
| value=DEFAULT_GPU_SECONDS, | |
| minimum=1, | |
| step=1, | |
| precision=0, | |
| interactive=True, | |
| scale=2, | |
| ) | |
| gr.Markdown("seconds", elem_classes=["unit-note"]) | |
| with gr.Row(): | |
| train_gpu_check = gr.Button("Gpu time check", scale=1) | |
| train_gpu_estimate = gr.Number( | |
| label="Estimated gpu time for these settings", | |
| value=None, | |
| precision=0, | |
| interactive=False, | |
| scale=2, | |
| ) | |
| with gr.Row(): | |
| push_to_hub = gr.Checkbox( | |
| label="Push Checkpoints to Hugging Face Hub", | |
| value=True, | |
| ) | |
| push_every_chunk = gr.Checkbox( | |
| label="Sync to Hub After Every Chunk (Crash-Proof Resume)", | |
| value=True, | |
| info="Uploads weights + optimizer.pt + ema.pt after each chunk.", | |
| ) | |
| resume_from_hub = gr.Checkbox( | |
| label="Resume from Hub Repo if Checkpoint Exists", | |
| value=True, | |
| info="Automatically pulls latest checkpoint from your HF repo before starting.", | |
| ) | |
| hf_private = gr.Checkbox( | |
| label="Private Hub Repository", | |
| value=True, | |
| ) | |
| with gr.Row(): | |
| hf_repo_id_input = gr.Textbox( | |
| label="Target Hugging Face Repo ID (Leave blank for `your-username/<lora-name>`)", | |
| placeholder="username/my-illustrious-lora", | |
| value="", | |
| ) | |
| # Hyperparameters & Model Options | |
| with gr.Accordion("🛠️ Training Hyperparameters & Memory Settings", open=False): | |
| with gr.Row(): | |
| lr = gr.Number(label="Base Learning Rate", value=1e-4, step=1e-5) | |
| unet_lr = gr.Number(label="UNet / Transformer LR", value=1e-4, step=1e-5) | |
| text_encoder_lr = gr.Number(label="Text Encoder LR", value=5e-5, step=1e-5) | |
| train_text_encoder = gr.Checkbox(label="Train Text Encoder (Recommended for Illustrious/SDXL)", value=True) | |
| with gr.Row(): | |
| rank = gr.Slider(label="LoRA Linear Rank (Dim)", minimum=4, maximum=128, step=4, value=16) | |
| alpha = gr.Slider(label="LoRA Linear Alpha", minimum=1, maximum=128, step=1, value=16) | |
| batch_size = gr.Slider(label="Batch Size", minimum=1, maximum=8, step=1, value=1) | |
| grad_accum = gr.Slider(label="Gradient Accumulation Steps", minimum=1, maximum=16, step=1, value=1) | |
| with gr.Row(): | |
| resolution_choice = gr.CheckboxGroup( | |
| choices=["512", "768", "1024"], | |
| value=["768", "1024"], | |
| label="Resolution Buckets", | |
| ) | |
| optimizer_type = gr.Dropdown( | |
| choices=["adamw8bit", "adamw", "prodigy", "dadaptation", "adafactor", "lion8bit"], | |
| value="adamw8bit", | |
| label="Optimizer", | |
| ) | |
| lr_scheduler_type = gr.Dropdown( | |
| choices=["cosine", "constant", "constant_with_warmup", "linear"], | |
| value="cosine", | |
| label="LR Scheduler", | |
| ) | |
| with gr.Row(): | |
| quantize_model = gr.Checkbox(label="Quantize Base Model (fp8 - Recommended for FLUX)", value=False) | |
| quantize_te = gr.Checkbox(label="Quantize Text Encoder (fp8)", value=False) | |
| low_vram = gr.Checkbox(label="Low VRAM Mode", value=False) | |
| use_ema = gr.Checkbox(label="Use EMA (Persisted across chunks)", value=True) | |
| with gr.Row(): | |
| cache_latents = gr.Checkbox(label="Cache Latents to Disk (Fastest chunk startup)", value=True) | |
| cache_text_embeddings = gr.Checkbox(label="Cache Text Embeddings to Disk (Disables TE training)", value=False) | |
| shuffle_tokens = gr.Checkbox(label="Shuffle Comma-Separated Tags (Recommended for Illustrious XL)", value=True) | |
| with gr.Accordion("Advanced YAML Config Overrides", open=False): | |
| use_advanced_yaml = gr.Checkbox(label="Enable YAML Overrides", value=False) | |
| advanced_yaml_str = gr.Code(value=DEFAULT_ADVANCED_YAML, language="yaml", label="process[0] overrides") | |
| # Validation / Sample Prompts | |
| with gr.Accordion("🖼️ Validation Sample Prompts", open=False, visible=False) as sample_accordion: | |
| with gr.Row(): | |
| disable_sampling = gr.Checkbox( | |
| label="Disable Sampling During Training (Saves GPU time per chunk)", | |
| value=False, | |
| ) | |
| sample_every = gr.Number(label="Sample Every N Steps", value=200, minimum=50, step=50) | |
| guidance_scale = gr.Number(label="Sample CFG Scale", value=6.0, minimum=1.0, maximum=15.0, step=0.5) | |
| sample_steps = gr.Number(label="Sample Inference Steps", value=28, minimum=4, maximum=60, step=1) | |
| sample_1 = gr.Textbox( | |
| label="Validation Prompt 1", | |
| value="1girl, solo, masterpiece, best quality, newest, detailed background", | |
| ) | |
| sample_2 = gr.Textbox( | |
| label="Validation Prompt 2", | |
| value="1girl, portrait, dramatic lighting, looking at viewer, masterpiece, best quality", | |
| ) | |
| sample_3 = gr.Textbox(label="Validation Prompt 3 (Optional)", value="") | |
| negative_prompt = gr.Textbox( | |
| label="Negative Prompt (For Illustrious XL / SDXL / SD1.5)", | |
| value=MODEL_PRESETS["Illustrious XL v0.1 (Diffusers Repo)"]["default_neg"], | |
| ) | |
| output_components.extend([sample_accordion, sample_1, sample_2, sample_3]) | |
| start_btn = gr.Button( | |
| "🚀 Start / Resume Chunked Training on ZeroGPU", | |
| variant="primary", | |
| size="lg", | |
| visible=True, | |
| ) | |
| output_components.append(start_btn) | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| training_log_md = gr.Markdown("### 📋 Training Status & Chunk Log\nUpload images and click **Start / Resume Chunked Training**.") | |
| latest_ckpt_file = gr.File(label="📥 Latest Trained LoRA Checkpoint (.safetensors)") | |
| with gr.Column(scale=2): | |
| sample_gallery = gr.Gallery(label="🎨 Generated Validation Samples", columns=2, height=380) | |
| # ================================================================== | |
| # TAB 2: TEST / INFERENCE WITH TRAINED LORA | |
| # ================================================================== | |
| with gr.Tab("🎨 Test Trained LoRA (ZeroGPU Inference)"): | |
| gr.Markdown( | |
| "### Test your trained **Illustrious XL**, **SDXL**, or **FLUX.1** LoRA immediately on ZeroGPU" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| inf_preset = gr.Dropdown( | |
| choices=list(MODEL_PRESETS.keys()), | |
| value="Illustrious XL v0.1 (Diffusers Repo)", | |
| label="Base Model Preset", | |
| ) | |
| inf_custom_base = gr.Textbox( | |
| label="Custom Base Model Repo / .safetensors (Optional)", | |
| placeholder="OnomaAIResearch/Illustrious-xl-early-release-v0", | |
| value="", | |
| ) | |
| inf_lora_path = gr.Textbox( | |
| label="LoRA Checkpoint Path or HF Repo ID", | |
| placeholder="output/my_illustrious_lora/my_illustrious_lora.safetensors or username/my-lora", | |
| value="output/my_illustrious_lora/my_illustrious_lora.safetensors", | |
| ) | |
| inf_prompt = gr.Textbox( | |
| label="Prompt", | |
| value="1girl, solo, masterpiece, best quality, newest, anime coloring, detailed background", | |
| lines=3, | |
| ) | |
| inf_neg_prompt = gr.Textbox( | |
| label="Negative Prompt", | |
| value=MODEL_PRESETS["Illustrious XL v0.1 (Diffusers Repo)"]["default_neg"], | |
| lines=2, | |
| ) | |
| with gr.Row(): | |
| inf_width = gr.Slider(label="Width", minimum=512, maximum=1536, step=64, value=1024) | |
| inf_height = gr.Slider(label="Height", minimum=512, maximum=1536, step=64, value=1024) | |
| with gr.Row(): | |
| inf_steps = gr.Slider(label="Steps", minimum=4, maximum=50, step=1, value=28) | |
| inf_cfg = gr.Slider(label="Guidance Scale (CFG)", minimum=1.0, maximum=12.0, step=0.5, value=6.0) | |
| with gr.Row(): | |
| inf_lora_scale = gr.Slider(label="LoRA Strength", minimum=0.0, maximum=2.0, step=0.05, value=1.0) | |
| inf_seed = gr.Number(label="Seed (-1 for random)", value=42, precision=0) | |
| with gr.Group(visible=True) as inf_gpu_group: | |
| gr.Markdown("**ZeroGPU**") | |
| with gr.Row(): | |
| inf_gpu_seconds = gr.Number( | |
| label="Gpu time to allocate", | |
| value=DEFAULT_INF_GPU_SECONDS, | |
| minimum=1, | |
| step=1, | |
| precision=0, | |
| interactive=True, | |
| scale=2, | |
| ) | |
| gr.Markdown("seconds", elem_classes=["unit-note"]) | |
| with gr.Row(): | |
| inf_gpu_check = gr.Button("Gpu time check", scale=1) | |
| inf_gpu_estimate = gr.Number( | |
| label="Estimated gpu time for these settings", | |
| value=None, | |
| precision=0, | |
| interactive=False, | |
| scale=2, | |
| ) | |
| inf_btn = gr.Button("✨ Generate Image on ZeroGPU", variant="primary") | |
| with gr.Column(scale=1): | |
| inf_output_img = gr.Image(label="Generated Image", type="pil") | |
| # ================================================================== | |
| # TAB 3: QUICK ZEROGPU TEST RUN (<= 4 MIN) | |
| # ================================================================== | |
| with gr.Tab("🧪 Quick ZeroGPU Test Run (Max 4 Min)"): | |
| gr.Markdown( | |
| "### Verify End-to-End Training on ZeroGPU (`@spaces.GPU(duration=120)`)\n" | |
| "Runs a self-contained synthetic training job (5 steps at 512x512, rank 4) within the 4-minute ZeroGPU limit " | |
| "to verify CUDA initialization, model loading, latent caching, LoRA backward pass, and `.safetensors` saving." | |
| ) | |
| with gr.Row(): | |
| smoke_preset = gr.Dropdown( | |
| choices=list(MODEL_PRESETS.keys()), | |
| value="Illustrious XL v0.1 (Diffusers Repo)", | |
| label="Model Preset to Test", | |
| ) | |
| smoke_steps = gr.Slider( | |
| label="Test Steps", | |
| minimum=2, | |
| maximum=20, | |
| step=1, | |
| value=5, | |
| ) | |
| smoke_gpu_seconds = gr.Number( | |
| label="Gpu time to allocate (seconds)", | |
| value=DEFAULT_GPU_SECONDS, | |
| minimum=1, | |
| step=1, | |
| precision=0, | |
| interactive=True, | |
| ) | |
| smoke_btn = gr.Button("🧪 Run 4-Min ZeroGPU Smoke Test", variant="primary") | |
| smoke_status_md = gr.Markdown("Click **Run 4-Min ZeroGPU Smoke Test** to verify.") | |
| smoke_ckpt_file = gr.File(label="Test LoRA Output (.safetensors)") | |
| # ====================================================================== | |
| # Wire Event Listeners | |
| # ====================================================================== | |
| preset_selector.change( | |
| fn=on_preset_change, | |
| inputs=[preset_selector], | |
| outputs=[ | |
| custom_model_path, | |
| custom_arch, | |
| lr, | |
| unet_lr, | |
| text_encoder_lr, | |
| train_text_encoder, | |
| quantize_model, | |
| quantize_te, | |
| resolution_choice, | |
| shuffle_tokens, | |
| guidance_scale, | |
| sample_steps, | |
| negative_prompt, | |
| preset_info_md, | |
| ], | |
| ) | |
| files_upload.upload( | |
| fn=handle_file_upload, | |
| inputs=[files_upload, trigger_word], | |
| outputs=output_components, | |
| ) | |
| do_captioning_btn.click( | |
| fn=handle_captioning_click, | |
| inputs=[uploaded_image_paths_state, trigger_word, caption_style, overwrite_captions] + caption_boxes, | |
| outputs=caption_boxes, | |
| ) | |
| train_gpu_check.click( | |
| fn=gpu_train_time_check, | |
| inputs=[ | |
| chunk_steps, | |
| resolution_choice, | |
| batch_size, | |
| grad_accum, | |
| preset_selector, | |
| train_text_encoder, | |
| disable_sampling, | |
| sample_steps, | |
| ], | |
| outputs=[train_gpu_estimate], | |
| ) | |
| inf_gpu_check.click( | |
| fn=gpu_inf_time_check, | |
| inputs=[inf_steps, inf_cfg, inf_width, inf_height], | |
| outputs=[inf_gpu_estimate], | |
| ) | |
| start_btn.click( | |
| fn=run_chunked_training_pipeline, | |
| inputs=[ | |
| uploaded_image_paths_state, | |
| lora_name, | |
| trigger_word, | |
| preset_selector, | |
| custom_model_path, | |
| custom_arch, | |
| total_steps, | |
| chunk_steps, | |
| auto_loop_chunks, | |
| lr, | |
| unet_lr, | |
| text_encoder_lr, | |
| train_text_encoder, | |
| rank, | |
| alpha, | |
| batch_size, | |
| grad_accum, | |
| resolution_choice, | |
| optimizer_type, | |
| lr_scheduler_type, | |
| quantize_model, | |
| quantize_te, | |
| low_vram, | |
| use_ema, | |
| cache_latents, | |
| cache_text_embeddings, | |
| shuffle_tokens, | |
| save_every, | |
| sample_every, | |
| disable_sampling, | |
| sample_1, | |
| sample_2, | |
| sample_3, | |
| negative_prompt, | |
| guidance_scale, | |
| sample_steps, | |
| push_to_hub, | |
| push_every_chunk, | |
| hf_repo_id_input, | |
| hf_private, | |
| resume_from_hub, | |
| hf_token_input, | |
| use_advanced_yaml, | |
| advanced_yaml_str, | |
| train_gpu_seconds, | |
| ] + caption_boxes, | |
| outputs=[training_log_md, latest_ckpt_file, sample_gallery], | |
| ) | |
| inf_btn.click( | |
| fn=gpu_test_inference, | |
| inputs=[ | |
| inf_preset, | |
| inf_custom_base, | |
| inf_lora_path, | |
| inf_prompt, | |
| inf_neg_prompt, | |
| inf_width, | |
| inf_height, | |
| inf_steps, | |
| inf_cfg, | |
| inf_lora_scale, | |
| inf_seed, | |
| hf_token_input, | |
| inf_gpu_seconds, | |
| ], | |
| outputs=[inf_output_img], | |
| ) | |
| smoke_btn.click( | |
| fn=gpu_quick_smoke_test, | |
| inputs=[smoke_preset, custom_model_path, custom_arch, smoke_steps, hf_token_input, smoke_gpu_seconds], | |
| outputs=[smoke_status_md, smoke_ckpt_file], | |
| ) | |
| return demo | |
| demo = build_ui() | |
| if __name__ == "__main__": | |
| demo.queue().launch(server_name="0.0.0.0", show_error=True) | |