344 lines
15 KiB
Python
344 lines
15 KiB
Python
# !/usr/bin/env python
|
|
# -*- coding: UTF-8 -*-
|
|
import sys
|
|
import os
|
|
import torch
|
|
from PIL import Image
|
|
import numpy as np
|
|
import random
|
|
from transformers import CLIPTokenizer, CLIPVisionModelWithProjection, CLIPImageProcessor
|
|
from diffusers import PriorTransformer, UNet2DConditionModel, KandinskyV22Pipeline
|
|
from .model import pops_utils
|
|
from .model.pipeline_pops import pOpsPipeline
|
|
import folder_paths
|
|
from comfy.utils import common_upscale
|
|
|
|
MAX_SEED = np.iinfo(np.int32).max
|
|
dir_path = os.path.dirname(os.path.abspath(__file__))
|
|
path_dir = os.path.dirname(dir_path)
|
|
file_path = os.path.dirname(path_dir)
|
|
|
|
device="cuda"
|
|
weight_dtype = torch.float16
|
|
output_dir = folder_paths.output_directory
|
|
|
|
def instance_path(path, repo):
|
|
if repo == "":
|
|
if path == "none":
|
|
raise "you need fill repo_id or download model in diffusers directory "
|
|
elif path != "none":
|
|
model_path = get_local_path(file_path, path)
|
|
repo = get_instance_path(model_path)
|
|
return repo
|
|
|
|
|
|
paths = []
|
|
for search_path in folder_paths.get_folder_paths("diffusers"):
|
|
if os.path.exists(search_path):
|
|
for root, subdir, files in os.walk(search_path, followlinks=True):
|
|
if "model_index.json" in files:
|
|
paths.append(os.path.relpath(root, start=search_path))
|
|
if paths:
|
|
paths = ["none"] + [x for x in paths if x]
|
|
else:
|
|
paths = ["none", ]
|
|
|
|
|
|
def phi2narry(img):
|
|
narry = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0)
|
|
return narry
|
|
|
|
def narry_list(list_in):
|
|
for i in range(len(list_in)):
|
|
value = list_in[i]
|
|
modified_value = phi2narry(value)
|
|
list_in[i] = modified_value
|
|
return list_in
|
|
|
|
def get_instance_path(path):
|
|
instance_path = os.path.normpath(path)
|
|
if sys.platform == 'win32':
|
|
instance_path = instance_path.replace('\\', "/")
|
|
return instance_path
|
|
|
|
|
|
def tensor_to_pil(tensor):
|
|
image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy()
|
|
image = Image.fromarray(image_np, mode='RGB')
|
|
return image
|
|
|
|
def nomarl_upscale_topil(img_tensor, width, height):
|
|
samples = img_tensor.movedim(-1, 1)
|
|
img = common_upscale(samples, width, height, "nearest-exact", "center")
|
|
samples = img.movedim(1, -1)
|
|
img_pil = tensor_to_pil(samples)
|
|
return img_pil
|
|
|
|
def get_local_path(file_path, model_path):
|
|
path = os.path.join(file_path, "models", "diffusers", model_path)
|
|
model_path = os.path.normpath(path)
|
|
if sys.platform == 'win32':
|
|
model_path = model_path.replace('\\', "/")
|
|
return model_path
|
|
|
|
|
|
def process_image(inputs_obj, clip,width, height):
|
|
image_caption_suffix = ''
|
|
if inputs_obj is not None and isinstance(inputs_obj, str):
|
|
if inputs_obj.suffix == '.pth':
|
|
image = torch.load(inputs_obj).image_embeds.to(device).to(weight_dtype)
|
|
image_caption_suffix = '(embedding)'
|
|
else:
|
|
raise "input embeds_a must be a pth file"
|
|
else:
|
|
if isinstance(inputs_obj, Image.Image):
|
|
image_pil =inputs_obj
|
|
else:
|
|
image_pil = Image.new('RGB', (height, width), (255, 255, 255))
|
|
image = torch.Tensor(clip(image_pil)['pixel_values'][0]).to(device).unsqueeze(0).to(weight_dtype)
|
|
return image, image_caption_suffix
|
|
|
|
|
|
def get_embedding(model,clip,inputs_obj, drop_condition_a, drop_condition_b,prior_seeds,prior_guidance_scale,prior_steps,height,width):
|
|
filename_prefix = ''.join(random.choice("0123456789") for _ in range(4))
|
|
should_drop_cond = [(drop_condition_a, drop_condition_b)]
|
|
image_list, image_pil_list, caption_suffix_list=[],[],[]
|
|
for input_path in inputs_obj:
|
|
# Process both inputs
|
|
image,caption_suffix = process_image(input_path, clip, width, height)
|
|
image_list.append(image)
|
|
caption_suffix_list.append(caption_suffix)
|
|
|
|
image_list=[i for i in image_list if i is not None]
|
|
input_image_embeds, input_hidden_state = pops_utils.preprocess(image_list[0], image_list[1],
|
|
model.image_encoder,
|
|
model.prior.clip_mean.detach(),
|
|
model.prior.clip_std.detach(),
|
|
should_drop_cond=should_drop_cond)
|
|
|
|
captions = [f"objects{caption_suffix_list[0]}", f"textures{caption_suffix_list[1]}"]
|
|
negative_input_embeds = torch.zeros_like(input_image_embeds)
|
|
negative_hidden_states = torch.zeros_like(input_hidden_state)
|
|
|
|
img_emb =model(input_embeds=input_image_embeds, input_hidden_states=input_hidden_state,
|
|
negative_input_embeds=negative_input_embeds,
|
|
negative_input_hidden_states=negative_hidden_states,
|
|
num_inference_steps=prior_steps,
|
|
num_images_per_prompt=1,
|
|
guidance_scale=prior_guidance_scale,
|
|
generator=torch.Generator(device=device).manual_seed(prior_seeds))
|
|
|
|
|
|
img_emb_file = os.path.join(output_dir, f"infer_binary{filename_prefix}_s_{prior_seeds}_cfg_{prior_guidance_scale}_img_emb.pth")
|
|
if not os.path.exists(img_emb_file):
|
|
torch.save(img_emb, img_emb_file)
|
|
positive_emb=img_emb.image_embeds
|
|
negative_emb = img_emb.negative_image_embeds
|
|
|
|
return positive_emb, negative_emb,input_hidden_state,img_emb_file
|
|
|
|
|
|
def get_embedding_instruct(model,clip,tokenizer, inputs_obj, texts,prior_guidance_scale, prior_seeds, prior_steps,height,width):
|
|
#print(texts, type(texts))
|
|
image_a, caption_suffix_a = process_image(inputs_obj,clip,width, height)
|
|
text_inputs = tokenizer(text=texts,padding="max_length",max_length=tokenizer.model_max_length,truncation=True,return_tensors="pt",)
|
|
mask = text_inputs.attention_mask.bool() # [0]
|
|
text_encoder_output = model.text_encoder(text_inputs.input_ids.to(device))
|
|
text_encoder_hidden_states = text_encoder_output.last_hidden_state
|
|
text_encoder_concat = text_encoder_hidden_states[:, :mask.sum().item()]
|
|
#
|
|
input_image_embeds, input_hidden_state = pops_utils.preprocess(image_a, None,
|
|
model.image_encoder,
|
|
model.prior.clip_mean.detach(),
|
|
model.prior.clip_std.detach(),
|
|
concat_hidden_states=text_encoder_concat)
|
|
|
|
negative_input_embeds = torch.zeros_like(input_image_embeds)
|
|
negative_hidden_states = torch.zeros_like(input_hidden_state)
|
|
# for scale in prior_guidance_scale:
|
|
img_emb = model(input_embeds=input_image_embeds, input_hidden_states=input_hidden_state,
|
|
negative_input_embeds=negative_input_embeds,
|
|
negative_input_hidden_states=negative_hidden_states,
|
|
num_inference_steps=prior_steps,
|
|
num_images_per_prompt=1,
|
|
guidance_scale=prior_guidance_scale,
|
|
generator=torch.Generator(device=device).manual_seed(prior_seeds))
|
|
|
|
img_emb_file = os.path.join(output_dir, f"{texts}_s_{prior_seeds}_cfg_{prior_guidance_scale}_img_emb.pth")
|
|
if not os.path.exists(img_emb_file):
|
|
torch.save(img_emb, img_emb_file)
|
|
positive_emb = img_emb.image_embeds
|
|
negative_emb = img_emb.negative_image_embeds
|
|
return positive_emb, negative_emb,input_hidden_state,img_emb_file
|
|
|
|
|
|
class Pops_Repo_Loader:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"local_prior": (paths,),
|
|
"prior_repo": ("STRING", {"default": "kandinsky-community/kandinsky-2-2-prior"}),
|
|
"local_decoder": (paths,),
|
|
"decoder_repo": ("STRING", {"default": "kandinsky-community/kandinsky-2-2-decoder"}),
|
|
"pops_ckpt": (["none"] + folder_paths.get_filename_list("checkpoints"),),
|
|
"function_type": (["Binary","instruct",],),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("MODEL","CLIP","VAE","MODEL",)
|
|
RETURN_NAMES = ("model","clip","vae","tokenizer",)
|
|
FUNCTION = "pops_repo_loader"
|
|
CATEGORY = "Pops"
|
|
|
|
def pops_repo_loader(self, local_prior, prior_repo, local_decoder, decoder_repo, pops_ckpt,function_type):
|
|
|
|
kandinsky_prior_repo = instance_path(local_prior, prior_repo)
|
|
kandinsky_decoder_repo = instance_path(local_decoder, decoder_repo)
|
|
Pops_ckpt = folder_paths.get_full_path("checkpoints", pops_ckpt)
|
|
#clip = CLIPImageProcessor.from_pretrained(kandinsky_prior_repo,subfolder='image_processor')
|
|
clip = CLIPImageProcessor()
|
|
prior = PriorTransformer.from_pretrained(kandinsky_prior_repo, subfolder="prior", use_safetensors=True)
|
|
prior_state_dict = torch.load(Pops_ckpt, map_location=torch.device('cuda'))
|
|
prior.load_state_dict(prior_state_dict, strict=False)
|
|
prior.eval()
|
|
prior = prior.to(weight_dtype)
|
|
image_encoder = CLIPVisionModelWithProjection.from_pretrained(kandinsky_prior_repo, subfolder='image_encoder',
|
|
torch_dtype=torch.float16).eval()
|
|
# Freeze text_encoder and image_encoder
|
|
image_encoder.requires_grad_(False)
|
|
|
|
model = pOpsPipeline.from_pretrained(kandinsky_prior_repo,
|
|
prior=prior,
|
|
image_encoder=image_encoder,
|
|
torch_dtype=torch.float16).to(device)
|
|
|
|
# Load decoder model f
|
|
unet = UNet2DConditionModel.from_pretrained(kandinsky_decoder_repo,
|
|
subfolder='unet').to(torch.float16).to(device)
|
|
vae = KandinskyV22Pipeline.from_pretrained(kandinsky_decoder_repo, unet=unet,
|
|
torch_dtype=torch.float16).to(device)
|
|
|
|
if function_type=="instruct":
|
|
tokenizer = CLIPTokenizer.from_pretrained(kandinsky_prior_repo, subfolder='tokenizer')
|
|
else:
|
|
tokenizer =None
|
|
return (model,clip,vae,tokenizer,)
|
|
|
|
|
|
class Pops_Sampler:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model":("MODEL",),
|
|
"clip": ("CLIP",),
|
|
"texts": ("STRING",{"default": "smooth"}),
|
|
"drop_condition_a": ("BOOLEAN", {"default": False},),
|
|
"drop_condition_b": ("BOOLEAN", {"default": False},),
|
|
"prior_guidance_scale": (
|
|
"FLOAT", {"default": 1.0, "min": 0.1, "max": 24.0, "step": 0.1, "round": False}),
|
|
"seed": ("INT", {"default": 2, "min": 1, "max": MAX_SEED}),
|
|
"prior_steps": ("INT", {"default": 25, "min": 1, "max": 4096}),
|
|
"height": ("INT", {"default": 768, "min": 256, "max": 4096,"step": 64}),
|
|
"width": ("INT", {"default": 768, "min": 256, "max": 4096,"step": 64}),
|
|
"use_mean": ("BOOLEAN", {"default": False},),},
|
|
"optional": {
|
|
"tokenizer": ("MODEL",),
|
|
"image_a": ("IMAGE",),
|
|
"image_b": ("IMAGE",),
|
|
"embeds_a": ("STRING", {"forceInput": True},),
|
|
"embeds_b": ("STRING", {"forceInput": True},),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("CONDITIONING", "CONDITIONING","STRING",)
|
|
RETURN_NAMES = ("positive_emb", "ng_image_embeds","img_emb_file",)
|
|
FUNCTION = "pops_sampler"
|
|
CATEGORY = "Pops"
|
|
|
|
def pops_sampler(self, model,clip,texts
|
|
, drop_condition_a,drop_condition_b,prior_guidance_scale,seed,prior_steps,height,width,use_mean,**kwargs):
|
|
tokenizer=kwargs.get("tokenizer")
|
|
image_a = kwargs.get("image_a")
|
|
image_b = kwargs.get("image_b")
|
|
if isinstance(image_a, torch.Tensor):
|
|
image_a = nomarl_upscale_topil(image_a, width, height)
|
|
if isinstance(image_b, torch.Tensor):
|
|
image_b = nomarl_upscale_topil(image_b, width, height)
|
|
embeds_a = kwargs.get("embeds_a")
|
|
embeds_b = kwargs.get("embeds_b")
|
|
inputs_obj=[image_a,image_b, embeds_a,embeds_b]
|
|
inputs_obj = [i for i in inputs_obj if i is not None]
|
|
if tokenizer:
|
|
inputs_obj=inputs_obj[0]
|
|
positive_emb, negative_emb,input_hidden_state,img_emb_file= get_embedding_instruct(model,tokenizer,clip, inputs_obj,texts,
|
|
prior_guidance_scale, seed, prior_steps,height,width)
|
|
else:
|
|
if len(inputs_obj)>2:
|
|
inputs_obj = random.choices(inputs_obj,k=2)
|
|
elif len(inputs_obj)<2:
|
|
inputs_obj=inputs_obj.append([None])
|
|
else:
|
|
pass
|
|
positive_emb, negative_emb,input_hidden_state,img_emb_file= get_embedding(model,clip,inputs_obj,drop_condition_a, drop_condition_b,seed,prior_guidance_scale,prior_steps,height,width)
|
|
if use_mean:
|
|
mean_emb = 0.5 * input_hidden_state[:, 0] + 0.5 * input_hidden_state[:, 1]
|
|
mean_emb = (mean_emb * model.prior.clip_std) + model.prior.clip_mean
|
|
negative_emb = model.get_zero_embed(mean_emb.shape[0], device=mean_emb.device)
|
|
positive_emb=mean_emb
|
|
return (positive_emb, negative_emb,img_emb_file,)
|
|
|
|
|
|
class Pops_Decoder:
|
|
def __init__(self):
|
|
pass
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"vae": ("VAE",),
|
|
"positive_emb": ("CONDITIONING",),
|
|
"negative_emb": ("CONDITIONING",),
|
|
"seed": ("INT", {"default": 2, "min": 1, "max": MAX_SEED}),
|
|
"steps": ("INT", {"default": 25, "min": 1, "max": 4096}),
|
|
"guidance_scale": (
|
|
"FLOAT", {"default": 1.0, "min": 0.1, "max": 24.0, "step": 0.1, "round": False}),
|
|
"height": ("INT", {"default": 768, "min": 256, "max": 4096, "step": 64}),
|
|
"width": ("INT", {"default": 768, "min": 256, "max": 4096, "step": 64}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("image",)
|
|
FUNCTION = "pops_decoder"
|
|
CATEGORY = "Pops"
|
|
|
|
def pops_decoder(self,vae,positive_emb,negative_emb,seed, steps,guidance_scale,height,width,):
|
|
images = vae(image_embeds=positive_emb, negative_image_embeds=negative_emb,
|
|
num_inference_steps=steps, height=height,
|
|
width=width, guidance_scale=guidance_scale,
|
|
generator=torch.Generator(device=device).manual_seed(seed)).images
|
|
images=phi2narry(images[0])
|
|
return (images,)
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"Pops_Repo_Loader": Pops_Repo_Loader,
|
|
"Pops_Sampler":Pops_Sampler,
|
|
"Pops_Decoder":Pops_Decoder,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"Pops_Repo_Loader": "Pops_Repo_Loader",
|
|
"Pops_Sampler":"Pops_Sampler",
|
|
"Pops_Decoder":"Pops_Decode",
|
|
}
|