|
|
|
|
|
import matplotlib.pyplot as plt |
|
import cv2 |
|
import torch |
|
import numpy as np |
|
from omegaconf import OmegaConf |
|
from PIL import Image |
|
from tqdm import tqdm, trange |
|
from imwatermark import WatermarkEncoder |
|
from itertools import islice |
|
import time |
|
from pytorch_lightning import seed_everything |
|
from torch import autocast |
|
import torchvision |
|
from ldm.util import instantiate_from_config |
|
from ldm.models.diffusion.ddim import DDIMSampler |
|
from torchvision.transforms import Resize |
|
import argparse |
|
import os |
|
import pathlib |
|
import glob |
|
import tqdm |
|
|
|
|
|
device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") |
|
|
|
def load_model_from_config(config): |
|
model = instantiate_from_config(config.model) |
|
|
|
model.cuda() |
|
model.eval() |
|
return model |
|
|
|
|
|
def get_model(config_path, ckpt_path, pretrained_sd_path): |
|
config = OmegaConf.load(f"{config_path}") |
|
model = load_model_from_config(config) |
|
model.load_state_dict(torch.load(pretrained_sd_path,map_location='cpu')['state_dict'],strict=False) |
|
|
|
|
|
pl_sd = torch.load(ckpt_path, map_location="cpu") |
|
wrapped_state_dict = pl_sd |
|
new_sd = {k.replace("_forward_module.", ""): wrapped_state_dict[k] for k in wrapped_state_dict} |
|
|
|
m, u = model.load_state_dict(new_sd, strict=False) |
|
if len(m) > 0: |
|
print("missing keys:") |
|
print(m) |
|
if len(u) > 0: |
|
print("unexpected keys:") |
|
print(u) |
|
|
|
|
|
model = model.to(device) |
|
return model |
|
|
|
def file_exists(path): |
|
""" Check if a file exists and is not a directory. """ |
|
if not os.path.isfile(path): |
|
raise argparse.ArgumentTypeError(f"{path} is not a valid file.") |
|
return path |
|
|
|
def parse_arguments(): |
|
""" Parses command-line arguments. """ |
|
parser = argparse.ArgumentParser(description="Process images based on provided paths.") |
|
parser.add_argument("--pretrained_sd", type=file_exists, required=True, help="Path to the SD1.4 pretrained checkpoint") |
|
parser.add_argument("--learned_params", type=file_exists, required=True, help="Path to the MagicFixup learned parameters.") |
|
parser.add_argument("--save_path", type=str, required=True, help="Path to save the full model state dict") |
|
|
|
return parser.parse_args() |
|
|
|
def main(): |
|
args = parse_arguments() |
|
model = get_model('configs/collage_mix_train.yaml',args.learned_params, args.pretrained_sd) |
|
torch.save(model.state_dict(), args.save_path) |
|
|
|
if __name__ == '__main__': |
|
main() |
