FillipandClaude Opus 4.6
8e16057d2d
Pixel-space diffusion text-to-image generation with REPA LoRA training. 6 nodes: Model Loader, Text Encoder Loader, Generate, Training Config, Train, LoRA Loader. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
128 lines
4.2 KiB
Python
128 lines
4.2 KiB
Python
"""Model loading and caching for PixelGen."""
|
|
|
|
import logging
|
|
import os
|
|
from typing import Optional
|
|
|
|
import torch
|
|
from huggingface_hub import hf_hub_download
|
|
|
|
from .model_info import MODEL_CONFIGS, TEXT_ENCODER_CONFIGS, AVAILABLE_MODELS
|
|
|
|
logger = logging.getLogger("FL_PixelGen")
|
|
|
|
LOADED_MODELS_CACHE = {}
|
|
LOADED_ENCODERS_CACHE = {}
|
|
|
|
DTYPE_MAP = {
|
|
"bfloat16": torch.bfloat16,
|
|
"float16": torch.float16,
|
|
"float32": torch.float32,
|
|
}
|
|
|
|
|
|
class PixelGenLoader:
|
|
@staticmethod
|
|
def load_model(model_name: str, device: str = "cuda", dtype: str = "bfloat16",
|
|
force_reload: bool = False):
|
|
cache_key = f"{model_name}_{device}_{dtype}"
|
|
if not force_reload and cache_key in LOADED_MODELS_CACHE:
|
|
logger.info(f"Using cached model: {model_name}")
|
|
return LOADED_MODELS_CACHE[cache_key]
|
|
|
|
from src.models.transformer.JiT_T2I import JiT_T2I
|
|
|
|
torch_dtype = DTYPE_MAP.get(dtype, torch.bfloat16)
|
|
|
|
# Resolve checkpoint path
|
|
if model_name in MODEL_CONFIGS:
|
|
config = MODEL_CONFIGS[model_name]
|
|
ckpt_path = PixelGenLoader._resolve_checkpoint(
|
|
config["repo_id"], config["filename"]
|
|
)
|
|
denoiser_config = config["denoiser_config"]
|
|
elif model_name in AVAILABLE_MODELS:
|
|
ckpt_path = AVAILABLE_MODELS[model_name]
|
|
denoiser_config = MODEL_CONFIGS.get(
|
|
"PixNerd-XXL-P16-T2I", {}
|
|
).get("denoiser_config", {})
|
|
else:
|
|
raise ValueError(f"Unknown model: {model_name}")
|
|
|
|
logger.info(f"Loading denoiser from: {ckpt_path}")
|
|
|
|
# Instantiate denoiser
|
|
denoiser = JiT_T2I(**denoiser_config)
|
|
|
|
# Load EMA weights
|
|
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
|
|
state_dict = ckpt.get("state_dict", ckpt)
|
|
prefix = "ema_denoiser."
|
|
loaded = 0
|
|
for k, v in denoiser.state_dict().items():
|
|
key = prefix + k
|
|
if key in state_dict:
|
|
v.copy_(state_dict[key])
|
|
loaded += 1
|
|
else:
|
|
logger.warning(f"Missing key in checkpoint: {key}")
|
|
logger.info(f"Loaded {loaded} weight tensors from checkpoint")
|
|
|
|
denoiser = denoiser.to(device=device, dtype=torch_dtype)
|
|
denoiser.eval()
|
|
|
|
result = {
|
|
"denoiser": denoiser,
|
|
"device": device,
|
|
"dtype": torch_dtype,
|
|
"model_name": model_name,
|
|
}
|
|
LOADED_MODELS_CACHE[cache_key] = result
|
|
return result
|
|
|
|
@staticmethod
|
|
def load_text_encoder(model_name: str, device: str = "cuda",
|
|
force_reload: bool = False):
|
|
cache_key = f"{model_name}_{device}"
|
|
if not force_reload and cache_key in LOADED_ENCODERS_CACHE:
|
|
logger.info(f"Using cached text encoder: {model_name}")
|
|
return LOADED_ENCODERS_CACHE[cache_key]
|
|
|
|
from src.models.conditioner.qwen3_text_encoder import Qwen3TextEncoder
|
|
|
|
config = TEXT_ENCODER_CONFIGS.get(model_name, {})
|
|
embed_dim = config.get("embed_dim", 2048)
|
|
max_length = config.get("max_length", 128)
|
|
|
|
logger.info(f"Loading text encoder: {model_name}")
|
|
conditioner = Qwen3TextEncoder(
|
|
weight_path=model_name,
|
|
embed_dim=embed_dim,
|
|
max_length=max_length,
|
|
)
|
|
conditioner = conditioner.to(device)
|
|
|
|
result = {
|
|
"conditioner": conditioner,
|
|
"device": device,
|
|
"model_name": model_name,
|
|
}
|
|
LOADED_ENCODERS_CACHE[cache_key] = result
|
|
return result
|
|
|
|
@staticmethod
|
|
def _resolve_checkpoint(repo_id: str, filename: str) -> str:
|
|
"""Download checkpoint from HF if not already local."""
|
|
# Check ComfyUI models directory first
|
|
import folder_paths
|
|
try:
|
|
local = folder_paths.get_full_path("pixelgen", filename)
|
|
if local is not None and os.path.isfile(local):
|
|
return local
|
|
except Exception:
|
|
pass
|
|
|
|
logger.info(f"Downloading {filename} from {repo_id}...")
|
|
path = hf_hub_download(repo_id=repo_id, filename=filename)
|
|
return path
|