# !/usr/bin/env python # -*- coding: UTF-8 -*- import os import numpy as np import torch from PIL import Image from omegaconf import OmegaConf from diffusers import ( AutoencoderKL, ControlNetModel, StableDiffusionXLControlNetPipeline, StableDiffusionXLPipeline, ) from transformers import BlipProcessor, BlipForConditionalGeneration from .CSGO.ip_adapter.utils import BLOCKS as BLOCKS from .CSGO.ip_adapter.utils import controlnet_BLOCKS as controlnet_BLOCKS from .CSGO.ip_adapter.utils import resize_content from .CSGO.ip_adapter import CSGO from safetensors.torch import load_file import folder_paths from comfy.utils import common_upscale from comfy.clip_vision import load MAX_SEED = np.iinfo(np.int32).max node_cur_path = os.path.dirname(os.path.abspath(__file__)) device = ( "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu" ) weight_dtype = torch.float16 def tensor2pil(tensor): image_np = tensor.squeeze().mul(255).clamp(0, 255).byte().numpy() image = Image.fromarray(image_np, mode='RGB') return image def pil2narry(img): narry = torch.from_numpy(np.array(img).astype(np.float32) / 255.0).unsqueeze(0) return narry def tensor_upscale2pil(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 = tensor2pil(samples) return img_pil def tensor_upscale(img_tensor, width, height): #torch tensor samples = img_tensor.movedim(-1, 1) img = common_upscale(samples, width, height, "nearest-exact", "center") samples = img.movedim(1, -1) return samples class Blip_Loader: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "blip_repo": ("STRING", {"default": "Salesforce/blip-image-captioning-large"}), } } RETURN_TYPES = ("MODEL", "MODEL",) RETURN_NAMES = ("blip_processor", "blip_model",) FUNCTION = "test" CATEGORY = "CSGO_Wrapper" def test(self, blip_repo,): blip_processor = BlipProcessor.from_pretrained(blip_repo) blip_model = BlipForConditionalGeneration.from_pretrained(blip_repo).to(device) return (blip_processor,blip_model,) class CSGO_Loader: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "base_cpkt":(["none"]+folder_paths.get_filename_list("checkpoints"),), "clip_vision":(["none"]+folder_paths.get_filename_list("clip_vision"),), "vae_id":(["none"]+folder_paths.get_filename_list("vae"),), "controlnet": (["none"]+folder_paths.get_filename_list("controlnet"),), "csgo_ckpt": (["none"]+folder_paths.get_filename_list("checkpoints"),), "num_content_tokens": ("INT", { "default": 4, "min": 1, # Minimum value "max": 512, # Maximum value "step": 1, # Slider's step "display": "number", # Cosmetic only: display as "number" or "slider" }), "num_style_tokens": ("INT", { "default": 32, "min": 1, # Minimum value "max": 1024, # Maximum value "step": 1, # Slider's step "display": "number", # Cosmetic only: display as "number" or "slider" }), } } RETURN_TYPES = ("MODEL",) RETURN_NAMES = ("csgo",) FUNCTION = "test" CATEGORY = "CSGO_Wrapper" def test(self,base_cpkt,clip_vision,vae_id,controlnet,csgo_ckpt,num_content_tokens,num_style_tokens): if csgo_ckpt=="none" or base_cpkt=="none" or controlnet=="none" or clip_vision=="none" : raise "need weight" csgo_ckpt=folder_paths.get_full_path("checkpoints",csgo_ckpt) ckpt_path = folder_paths.get_full_path("checkpoints", base_cpkt) clip_vision=folder_paths.get_full_path("clip_vision", clip_vision) controlnet_path=folder_paths.get_full_path("controlnet", controlnet) img_encoder=load(clip_vision) model_config = os.path.join(node_cur_path, "local_repo") original_config_file = os.path.join(node_cur_path, 'configs', 'sd_xl_base.yaml') try: pipe = StableDiffusionXLPipeline.from_single_file( ckpt_path, config=model_config, original_config=original_config_file, torch_dtype=torch.float16) except: try: pipe = StableDiffusionXLPipeline.from_single_file( ckpt_path, config=model_config, original_config_file=original_config_file, torch_dtype=torch.float16) except: raise "load pipe error!,check you diffusers" controlnet = ControlNetModel.from_unet(pipe.unet) cn_state_dict = load_file(controlnet_path, device="cpu") controlnet.load_state_dict(cn_state_dict, strict=False) controlnet.to(torch.float16) try: pipe = StableDiffusionXLControlNetPipeline.from_single_file( ckpt_path,unet=pipe.unet,config=model_config, original_config=original_config_file, controlnet=controlnet, torch_dtype=torch.float16) except: try: pipe = StableDiffusionXLControlNetPipeline.from_single_file( ckpt_path,unet=pipe.unet,config=model_config, original_config_file=original_config_file,controlnet=controlnet, torch_dtype=torch.float16) except: raise "load pipe error!,check you diffusers" if vae_id!="none": vae_id = folder_paths.get_full_path("vae", vae_id) vae_config = os.path.join(node_cur_path, "local_repo", "vae") pipe.vae = AutoencoderKL.from_single_file(vae_id, config=vae_config, torch_dtype=torch.float16) pipe.enable_vae_tiling() pipe.enable_freeu(s1=0.6, s2=0.4, b1=1.1, b2=1.2) if device != "mps": pipe.enable_model_cpu_offload() torch.cuda.empty_cache() target_content_blocks = BLOCKS['content'] target_style_blocks = BLOCKS['style'] controlnet_target_content_blocks = controlnet_BLOCKS['content'] controlnet_target_style_blocks = controlnet_BLOCKS['style'] config_path = os.path.join(node_cur_path, "configs", "config.json") image_encoder_config = OmegaConf.load(config_path) csgo = CSGO(pipe, img_encoder, csgo_ckpt, device, image_encoder_config,num_content_tokens=num_content_tokens, num_style_tokens=num_style_tokens, target_content_blocks=target_content_blocks, target_style_blocks=target_style_blocks, controlnet_adapter=True, controlnet_target_content_blocks=controlnet_target_content_blocks, controlnet_target_style_blocks=controlnet_target_style_blocks, content_model_resampler=True, style_model_resampler=True, ) return (csgo,) class CSGO_Sampler: def __init__(self): pass @classmethod def INPUT_TYPES(s): return { "required": { "content_image": ("IMAGE",), "style_image": ("IMAGE",), "csgo": ("MODEL",), "prompt": ("STRING", {"multiline": True,"default": "a cat"}), "negative_prompt": ("STRING", {"multiline": True,"default": "text, watermark, lowres, low quality, worst quality, deformed, glitch, low contrast, noisy, saturation, blurry"}), "num_sample": ("INT", { "default": 1, "min": 1, # Minimum value "max": 100, # Maximum value "step": 1, # Slider's step "display": "number", # Cosmetic only: display as "number" or "slider" }), "width": ("INT", { "default": 1024, "min": 256, # Minimum value "max": 4096, # Maximum value "step": 64, # Slider's step "display": "number", # Cosmetic only: display as "number" or "slider" }), "height": ("INT", { "default": 1024, "min": 256, # Minimum value "max": 4096, # Maximum value "step": 64, # Slider's step "display": "number", # Cosmetic only: display as "number" or "slider" }), "content_scale": ("FLOAT", { "default": 0.5, "min": 0.1, "max": 1.0, "step": 0.01, "round": 0.001, # The value representing the precision to round to, will be set to the step value by default. Can be set to False to disable rounding. "display": "number", }), "style_scale": ("FLOAT", { "default": 1.0, "min": 0.1, "max": 10.0, "step": 0.01, "round": 0.001, # The value representing the precision to round to, will be set to the step value by default. Can be set to False to disable rounding. "display": "number", }), "cfg": ("FLOAT", { "default": 10.0, "min": 0.1, "max": 30.0, "step": 0.01, "round": 0.001, # The value representing the precision to round to, will be set to the step value by default. Can be set to False to disable rounding. "display": "number", }), "steps": ("INT", { "default": 25, "min": 1, # Minimum value "max": 100, # Maximum value "step": 1, # Slider's step "display": "number", # Cosmetic only: display as "number" or "slider" }), "seed": ("INT", { "default": 42, "min": -1, # Minimum value "max": MAX_SEED, # Maximum value }), "controlnet_scale": ("FLOAT", { "default": 0.6, "min": 0.1, "max": 1.0, "step": 0.01, "round": 0.001, # The value representing the precision to round to, will be set to the step value by default. Can be set to False to disable rounding. "display": "number", }), "text_only": ("BOOLEAN", {"default": False},), }, "optional": { "blip_processor": ("MODEL",), "blip_model": ("MODEL",), }, } RETURN_TYPES = ("IMAGE",) RETURN_NAMES = ("images",) FUNCTION = "test" CATEGORY = "CSGO_Wrapper" def test(self,content_image,style_image, csgo, prompt,negative_prompt,num_sample, width,height,content_scale,style_scale,cfg,steps,seed,controlnet_scale,text_only,**kwargs): blip_processor=kwargs.get("blip_processor",None) blip_model=kwargs.get("blip_model",None) #向量转图片 #style_image= tensor2pil(style_image) #content_image= tensor2pil(content_image) content_image_pil = tensor_upscale2pil(content_image, width, height) style_image =tensor_upscale(style_image, width, height) # torch tensor content_image=tensor_upscale(content_image, width, height) # torch tensor #clip vison if blip_processor and blip_model: #判断模型是否存在，不存在则跑文本驱动流程 with torch.no_grad(): inputs = blip_processor(content_image, return_tensors="pt").to(device) out = blip_model.generate(**inputs) caption = blip_processor.decode(out[0], skip_special_tokens=True) del blip_model,blip_processor torch.cuda.empty_cache() else: caption = prompt if text_only: content_image_pil = Image.fromarray( np.zeros((content_image_pil.size[0], content_image_pil.size[1], 3), dtype=np.uint8)).convert('RGB') content_image=pil2narry(content_image_pil) #width, height, content_image = resize_content(content_image)#? images = csgo.generate(pil_content_image=content_image, pil_style_image=style_image, prompt=caption, negative_prompt=negative_prompt, height=height, width=width, content_scale=content_scale, style_scale=style_scale, guidance_scale=cfg, num_images_per_prompt=num_sample, num_samples=num_sample, num_inference_steps=steps, seed=seed, image= content_image_pil, controlnet_conditioning_scale=controlnet_scale, )[0] images=pil2narry(images) return (images,) #@classmethod #def IS_CHANGED(s, image, string_field, int_field, float_field, print_to_screen): # return "" # Set the web directory, any .js file in that directory will be loaded by the frontend as a frontend extension # WEB_DIRECTORY = "./somejs" # A dictionary that contains all nodes you want to export with their names # NOTE: names should be globally unique NODE_CLASS_MAPPINGS = { "Blip_Loader":Blip_Loader, "CSGO_Loader": CSGO_Loader, "CSGO_Sampler":CSGO_Sampler } # A dictionary that contains the friendly/humanly readable titles for the nodes NODE_DISPLAY_NAME_MAPPINGS = { "Blip_Loader":"Blip_Loader", "CSGO_Loader": "CSGO_Loader", "CSGO_Sampler":"CSGO_Sampler" }