# /// script # requires-python = "==3.11.*" # dependencies = [ # "torch==2.7.1", "torchvision==0.22.1", "diffusers==0.25.0", # "transformers==4.28.1", "peft==0.9.0", "accelerate==0.27.2", # "huggingface-hub==0.25.2", "numpy==1.26.4", "onnx==1.19.0", # ] # [tool.uv.sources] # torch = { index = "pytorch" } # torchvision = { index = "pytorch" } # [[tool.uv.index]] # name = "pytorch" # url = "https://download.pytorch.org/whl/cu128" # explicit = true # /// """Export PiSA-SR components from pinned upstream weights on CUDA.""" import argparse import hashlib import json import sys from pathlib import Path SOURCE_COMMIT = "9c76e12e261d1980148738a0cff41e5e27d18c9f" BASE_REPOSITORY = "sd2-community/stable-diffusion-2-1-base" BASE_REVISION = "4e63672c03103b6c636b8fb4119ba982469b2955" CHECKPOINT_SHA256 = "13502f7bde5e4e585f84f24c3d904d9c848e068bf89f984046fe87da81a391f6" CHECKPOINT_URL = "https://drive.google.com/uc?export=download&id=1PWFM5zUd6lNMOyKX64anvUXhKEOCrf8L" def finalize_assets(destination): import shutil import onnx for path in destination.glob("*.onnx"): model = onnx.load(path) onnx.helper.set_model_props(model, { "license": "Apache-2.0; CreativeML Open RAIL++-M (LICENSE-MODEL.txt)", "modifications": "AnyImage: merged PiSA-SR LoRA; FP16; fixed empty prompt and timestep; split VAE. See NOTICE.txt.", "source_commit": SOURCE_COMMIT, }) onnx.save(model, path) for path in (Path(__file__).parent / "licenses/pisa-sr").glob("*.txt"): shutil.copyfile(path, destination / path.name) def prepare_sources(cache): import io import zipfile from urllib.request import urlopen from huggingface_hub import snapshot_download source = cache / f"PiSA-SR-{SOURCE_COMMIT}" if not source.exists(): url = f"https://codeload.github.com/csslc/PiSA-SR/zip/{SOURCE_COMMIT}" with urlopen(url, timeout=60) as response: archive = zipfile.ZipFile(io.BytesIO(response.read())) for member in archive.infolist(): destination = (cache / member.filename).resolve() if not destination.is_relative_to(cache.resolve()): raise ValueError("Invalid upstream archive path") archive.extractall(cache) checkpoint = cache / "pisa_sr.pkl" if not checkpoint.exists(): with urlopen(CHECKPOINT_URL, timeout=60) as response: checkpoint.write_bytes(response.read()) if hashlib.sha256(checkpoint.read_bytes()).hexdigest() != CHECKPOINT_SHA256: raise ValueError("PiSA-SR checkpoint checksum mismatch") base = snapshot_download( BASE_REPOSITORY, revision=BASE_REVISION, local_dir=cache / "base-model", allow_patterns=["*/*.json", "tokenizer/*", "*/*.fp16.safetensors"], ) return source, Path(base), checkpoint def export_components(source, base, checkpoint, destination): import gc import numpy as np import torch from peft import LoraConfig from diffusers.utils import set_weights_and_activate_adapters from transformers import AutoTokenizer, CLIPTextModel if hashlib.sha256(checkpoint.read_bytes()).hexdigest() != CHECKPOINT_SHA256: raise ValueError("PiSA-SR checkpoint checksum mismatch") sys.path.insert(0, str(source)) from src.models.autoencoder_kl import AutoencoderKL from src.models.unet_2d_condition import UNet2DConditionModel class Residual(torch.nn.Module): def __init__(self, model, prompt): super().__init__() self.model = model self.register_buffer("prompt", prompt) self.register_buffer("timestep", torch.tensor([1], dtype=torch.long)) def forward(self, latent): return self.model(latent, self.timestep, encoder_hidden_states=self.prompt).sample class Encoder(torch.nn.Module): def __init__(self, model): super().__init__() self.encoder = model.encoder self.quant_conv = model.quant_conv def forward(self, image): return self.quant_conv(self.encoder(image)) class Decoder(torch.nn.Module): def __init__(self, model): super().__init__() self.decoder = model.decoder self.post_quant_conv = model.post_quant_conv def forward(self, latent): return self.decoder(self.post_quant_conv(latent / 0.18215)).clamp(-1, 1) def export(wrapper, component, tensor): wrapper.eval().cuda() with torch.inference_mode(): torch.onnx.export( wrapper, (tensor.cuda(),), str(destination / f"{component}.onnx"), opset_version=17, input_names=["input"], output_names=["output"], do_constant_folding=True, dynamo=False, ) wrapper.cpu() torch.cuda.empty_cache() destination.mkdir(parents=True, exist_ok=True) torch.set_num_threads(8) torch.backends.cuda.matmul.allow_tf32 = False tokenizer = AutoTokenizer.from_pretrained(base, subfolder="tokenizer") text = CLIPTextModel.from_pretrained( base, subfolder="text_encoder", variant="fp16", torch_dtype=torch.float16, ).eval().cuda() with torch.inference_mode(): tokens = tokenizer("", max_length=tokenizer.model_max_length, padding="max_length", truncation=True, return_tensors="pt") prompt = text(tokens.input_ids.cuda())[0].cpu() del text torch.cuda.empty_cache() state = torch.load(checkpoint, map_location="cpu", weights_only=True) latent = torch.zeros((1, 4, 64, 64), dtype=torch.float16) unet = UNet2DConditionModel.from_pretrained(base, subfolder="unet", variant="fp16") for strength, component in (("pix", "unet-pixel"), ("sem", "unet-full")): adapters = [] for part in ("encoder", "decoder", "others"): key = f"{part}_modules" config = LoraConfig( r=state[f"lora_rank_unet_{strength}"], init_lora_weights="gaussian", target_modules=state[f"unet_lora_{key}_{strength}"], ) name = f"default_{part}_{strength}" unet.add_adapter(config, adapter_name=name) adapters.append(name) for name, parameter in unet.named_parameters(): if "lora" in name: parameter.data.copy_(state["state_dict_unet"][name]) set_weights_and_activate_adapters(unet, adapters, [1.0]*3) unet.merge_and_unload() # Export a half-precision copy; retain FP32 weights for the semantic merge. import copy wrapper = Residual(copy.deepcopy(unet).half(), prompt) export(wrapper, component, latent) del wrapper gc.collect() del unet, state gc.collect() vae = AutoencoderKL.from_pretrained(base, subfolder="vae", variant="fp16", torch_dtype=torch.float16) export(Encoder(vae), "encoder", torch.zeros((1,3,512,512), dtype=torch.float16)) export(Decoder(vae), "decoder", latent) torch.manual_seed(42) np.save(destination / "noise.npy", torch.randn((1,4,64,64), device="cuda", dtype=torch.float16).cpu().numpy()) finalize_assets(destination) files = [] for path in sorted(destination.iterdir()): if path.suffix in (".onnx", ".npy", ".txt"): with path.open("rb") as stream: checksum = hashlib.file_digest(stream, "sha256").hexdigest() files.append((path.name, path.stat().st_size, checksum)) (destination / "provenance.json").write_text(json.dumps({ "source_commit": SOURCE_COMMIT, "base_repository": BASE_REPOSITORY, "base_revision": BASE_REVISION, "checkpoint_sha256": CHECKPOINT_SHA256, "licenses": ["Apache-2.0", "CreativeML Open RAIL++-M"], "files": files, }, indent=2)) def main(): parser = argparse.ArgumentParser() parser.add_argument("--cache", type=Path, default=Path(".model-cache/pisa-sr")) parser.add_argument("--destination", type=Path, default=Path("models/pisa-sr")) args = parser.parse_args() args.cache.mkdir(parents=True, exist_ok=True) export_components(*prepare_sources(args.cache), args.destination) if __name__ == "__main__": main()