github/ai-toolkit
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Save entire pixart model again

Browse Source
This commit is contained in:
Jaret Burkett
2024-07-07 07:56:48 -06:00
parent cab8a1c7b8
commit 045e4a6e15
5 changed files with 177 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
testing/merge_in_text_encoder_adapter.py
View File

@@ -2,7 +2,7 @@ import os
import torch
from transformers import T5EncoderModel, T5Tokenizer
from diffusers import StableDiffusionPipeline, UNet2DConditionModel, PixArtSigmaPipeline, Transformer2DModel
from diffusers import StableDiffusionPipeline, UNet2DConditionModel, PixArtSigmaPipeline, Transformer2DModel, PixArtTransformer2DModel
from safetensors.torch import load_file, save_file
from collections import OrderedDict
import json
@@ -13,8 +13,8 @@ import json
# output_path = "/home/jaret/Dev/models/hf/kl-f16-d42_sd15_t5xl_raw"
model_path = "/home/jaret/Dev/models/hf/PixArt-Sigma-XL-2-1024-MS"
te_path = "google/flan-t5-base"
te_aug_path = "/home/jaret/Dev/models/tmp/pixart_sigma_t5base_000204000.safetensors"
output_path = "/home/jaret/Dev/models/hf/PixArt-Sigma-XL-2-1024-MS_t5base_raw"
te_aug_path = "/home/jaret/Dev/models/tmp/pixart_sigma_t5base_000227500.safetensors"
output_path = "/home/jaret/Dev/models/hf/PixArt-Sigma-XL-2-512_MS_t5base_raw"
print("Loading te adapter")
@@ -28,11 +28,13 @@ is_pixart = "pixart" in model_path.lower()
pipeline_class = StableDiffusionPipeline
transformer = PixArtTransformer2DModel.from_pretrained('PixArt-alpha/PixArt-Sigma-XL-2-512-MS', subfolder='transformer', torch_dtype=torch.float16)
if is_pixart:
pipeline_class = PixArtSigmaPipeline
if is_diffusers:
sd = pipeline_class.from_pretrained(model_path, torch_dtype=torch.float16)
sd = pipeline_class.from_pretrained(model_path, transformer=transformer, torch_dtype=torch.float16)
else:
sd = pipeline_class.from_single_file(model_path, torch_dtype=torch.float16)
@@ -113,6 +115,11 @@ for name in attn_processor_keys:
if weight_idx > 1000:
raise ValueError("Could not find the next weight")
orig_weight_shape_k = list(unet_sd[layer_name + ".to_k.weight"].shape)
new_weight_shape_k = list(te_aug_sd[te_aug_name + ".weight"].shape)
orig_weight_shape_v = list(unet_sd[layer_name + ".to_v.weight"].shape)
new_weight_shape_v = list(te_aug_sd[te_aug_name.replace('to_k', 'to_v') + ".weight"].shape)
unet_sd[layer_name + ".to_k.weight"] = te_aug_sd[te_aug_name + ".weight"]
unet_sd[layer_name + ".to_v.weight"] = te_aug_sd[te_aug_name.replace('to_k', 'to_v') + ".weight"]
@@ -120,6 +127,14 @@ for name in attn_processor_keys:
new_cross_attn_dim = unet_sd[layer_name + ".to_k.weight"].shape[1]
if is_pixart:
# copy the caption_projection weight
del unet_sd['caption_projection.linear_1.bias']
del unet_sd['caption_projection.linear_1.weight']
del unet_sd['caption_projection.linear_2.bias']
del unet_sd['caption_projection.linear_2.weight']
print("Saving unmodified model")
sd = sd.to("cpu", torch.float16)
sd.save_pretrained(
@@ -150,7 +165,7 @@ with open(os.path.join(unet_folder, "config.json"), 'r') as f:
config['cross_attention_dim'] = new_cross_attn_dim
if is_pixart:
config['caption_channels'] = te.config.d_model
config['caption_channels'] = None
# save it
with open(os.path.join(unet_folder, "config.json"), 'w') as f: