diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..855028e --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ + +from .csgo_node import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/configs/sd_xl_base.yaml b/configs/sd_xl_base.yaml new file mode 100644 index 0000000..6047379 --- /dev/null +++ b/configs/sd_xl_base.yaml @@ -0,0 +1,93 @@ +model: + target: sgm.models.diffusion.DiffusionEngine + params: + scale_factor: 0.13025 + disable_first_stage_autocast: True + + denoiser_config: + target: sgm.modules.diffusionmodules.denoiser.DiscreteDenoiser + params: + num_idx: 1000 + + scaling_config: + target: sgm.modules.diffusionmodules.denoiser_scaling.EpsScaling + discretization_config: + target: sgm.modules.diffusionmodules.discretizer.LegacyDDPMDiscretization + + network_config: + target: sgm.modules.diffusionmodules.openaimodel.UNetModel + params: + adm_in_channels: 2816 + num_classes: sequential + use_checkpoint: True + in_channels: 4 + out_channels: 4 + model_channels: 320 + attention_resolutions: [4, 2] + num_res_blocks: 2 + channel_mult: [1, 2, 4] + num_head_channels: 64 + use_linear_in_transformer: True + transformer_depth: [1, 2, 10] + context_dim: 2048 + spatial_transformer_attn_type: softmax-xformers + + conditioner_config: + target: sgm.modules.GeneralConditioner + params: + emb_models: + - is_trainable: False + input_key: txt + target: sgm.modules.encoders.modules.FrozenCLIPEmbedder + params: + layer: hidden + layer_idx: 11 + + - is_trainable: False + input_key: txt + target: sgm.modules.encoders.modules.FrozenOpenCLIPEmbedder2 + params: + arch: ViT-bigG-14 + version: laion2b_s39b_b160k + freeze: True + layer: penultimate + always_return_pooled: True + legacy: False + + - is_trainable: False + input_key: original_size_as_tuple + target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND + params: + outdim: 256 + + - is_trainable: False + input_key: crop_coords_top_left + target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND + params: + outdim: 256 + + - is_trainable: False + input_key: target_size_as_tuple + target: sgm.modules.encoders.modules.ConcatTimestepEmbedderND + params: + outdim: 256 + + first_stage_config: + target: sgm.models.autoencoder.AutoencoderKL + params: + embed_dim: 4 + monitor: val/rec_loss + ddconfig: + attn_type: vanilla-xformers + double_z: true + z_channels: 4 + resolution: 256 + in_channels: 3 + out_ch: 3 + ch: 128 + ch_mult: [1, 2, 4, 4] + num_res_blocks: 2 + attn_resolutions: [] + dropout: 0.0 + lossconfig: + target: torch.nn.Identity diff --git a/csgo_node.py b/csgo_node.py new file mode 100644 index 0000000..8069aa6 --- /dev/null +++ b/csgo_node.py @@ -0,0 +1,325 @@ +# !/usr/bin/env python +# -*- coding: UTF-8 -*- +import os +import numpy as np +import torch +import cv2 +from PIL import Image + +from transformers import AutoImageProcessor, AutoModel +from diffusers import ( + AutoencoderKL, + ControlNetModel, + StableDiffusionXLControlNetPipeline, + +) +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 +MAX_SEED = np.iinfo(np.int32).max +import folder_paths +from comfy.utils import common_upscale + +node_cur_path = os.path.dirname(os.path.abspath(__file__)) +#device = torch.device("cuda" if torch.cuda.is_available() else "cpu") #cuda:0 +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_Example_Test" + + 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"),), + "image_encoder_repo": ("STRING", {"default": "h94/IP-Adapter"}), + "vae_id":(["none"]+folder_paths.get_filename_list("vae"),), + "controlnet_repo": ("STRING", {"default": "TTPlanet/TTPLanet_SDXL_Controlnet_Tile_Realistic" }), + "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_Example_Test" + + def test(self,base_cpkt,image_encoder_repo,vae_id,controlnet_repo,csgo_ckpt,num_content_tokens,num_style_tokens): + + csgo_ckpt=folder_paths.get_full_path("checkpoints",csgo_ckpt) #获取绝对路径 + ckpt_path = folder_paths.get_full_path("checkpoints", base_cpkt) # 获取绝对路径 + + #vae = AutoencoderKL.from_pretrained(vae, torch_dtype=torch.float16) + controlnet = ControlNetModel.from_pretrained(controlnet_repo, torch_dtype=torch.float16, use_safetensors=True) + # pipe = StableDiffusionXLControlNetPipeline.from_pretrained( + # base_repo, + # controlnet=controlnet, + # torch_dtype=torch.float16, + # add_watermarker=False, + # vae=vae + # ) + + vae_id = folder_paths.get_full_path("vae", vae_id) + vae = AutoencoderKL.from_single_file(vae_id, torch_dtype=torch.float16) + + original_config_file = os.path.join(node_cur_path, 'configs', 'sd_xl_base.yaml') + try: + pipe = StableDiffusionXLControlNetPipeline.from_single_file( + ckpt_path, original_config=original_config_file,vae=vae, controlnet=controlnet, + torch_dtype=torch.float16) + except: + try: + pipe = StableDiffusionXLControlNetPipeline.from_single_file( + ckpt_path, original_config_file=original_config_file,vae=vae, controlnet=controlnet, + torch_dtype=torch.float16) + except: + raise "load pipe error!,check you diffusers" + + pipe.enable_vae_tiling() + #pipe.to(device) + 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'] + + csgo = CSGO(pipe, image_encoder_repo, csgo_ckpt, device, 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": 50, + "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_Example_Test" + + 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 = tensor_upscale2pil(content_image, width, height) + #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 = Image.fromarray( + np.zeros((content_image.size[0], content_image.size[1], 3), dtype=np.uint8)).convert('RGB') + + #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.convert('RGB'), + 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" +} diff --git a/example/workflow.png b/example/workflow.png new file mode 100644 index 0000000..10c16e8 Binary files /dev/null and b/example/workflow.png differ