nan
89087fc2e2
- Implemented the "O Image Relief Plane" Geometry Nodes group to create a fixed-base image relief using depth images and various parameters. - Created the "O Shadeless" shader node group to display color without scene lighting while preserving shadows and reflections. - Introduced a preview feature for node groups, allowing for interactive visualization of node arrangements in an HTML format. - Added a template HTML file for rendering node previews with customizable views and controls.
199 lines
8.1 KiB
Python
199 lines
8.1 KiB
Python
# /// 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()
|