feat: Add new Geometry Nodes groups for image relief and shadeless shader
- 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.
This commit is contained in:
198
tools/models/pisa_sr_export.py
Normal file
198
tools/models/pisa_sr_export.py
Normal file
@@ -0,0 +1,198 @@
|
||||
# /// 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()
|
||||
Reference in New Issue
Block a user