diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..1bd7413 --- /dev/null +++ b/__init__.py @@ -0,0 +1,4 @@ + +from .nodes import NODE_CLASS_MAPPINGS +NODE_DISPLAY_NAME_MAPPINGS = {k:v.TITLE for k,v in NODE_CLASS_MAPPINGS.items()} +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..89aa7db --- /dev/null +++ b/nodes.py @@ -0,0 +1,91 @@ +import sys +import os +import importlib + +# sys.path.insert(0, os.path.dirname(os.path.realpath(__file__))) + +class Kiki_Omini_Subject: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "subject_image": ("IMAGE",), + "prompt": ("STRING", {"multiline": True, + "default": ''}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, + "tooltip": "The random seed used for creating the noise."}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "run" + TITLE = 'OminiControl Subject' + + CATEGORY = "Runninghub/Omini" + DESCRIPTION = "Ominicontrol subject node" + + def run(self, subject_image, prompt, seed): + import ComfyUI_RH_OminiControl.rh_omini_subject as ros + img = ros.run(subject_image, prompt, seed) + return (img, ) + +class Kiki_Omini_Spatial: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "ref_image": ("IMAGE",), + "prompt": ("STRING", {"multiline": True, + "default": ''}), + "condition_type": (["canny", "depth", "coloring", "deblurring"], ), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, + "tooltip": "The random seed used for creating the noise."}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "run" + TITLE = 'OminiControl Spatial' + + CATEGORY = "Runninghub/Omini" + DESCRIPTION = "Ominicontrol spatial node" + + def run(self, ref_image, prompt, condition_type, seed): + import ComfyUI_RH_OminiControl.rh_omini_spatial as rosp + img = rosp.run(ref_image, prompt, condition_type, seed) + return (img, ) + +class Kiki_Omini_Fill: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "ori_image": ("IMAGE",), + "mask": ("MASK", ), + "prompt": ("STRING", {"multiline": True, + "default": ''}), + "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff, + "tooltip": "The random seed used for creating the noise."}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "run" + TITLE = 'OminiControl Fill' + + CATEGORY = "Runninghub/Omini" + DESCRIPTION = "Ominicontrol fill node" + + def run(self, ori_image, mask, prompt, seed): + import ComfyUI_RH_OminiControl.rh_omini_fill as rof + img = rof.run(ori_image, mask, prompt, seed) + return (img, ) + +NODE_CLASS_MAPPINGS = { + "RunningHub_Omini_Subject": Kiki_Omini_Subject, + "RunningHub_Omini_Spatial": Kiki_Omini_Spatial, + "RunningHub_Omini_Fill": Kiki_Omini_Fill, +} diff --git a/readme b/readme new file mode 100644 index 0000000..c4225af --- /dev/null +++ b/readme @@ -0,0 +1,16 @@ +原项目地址：https://github.com/Yuanshi9815/OminiControl + +通过分拆pipeline加载，让该项目可以在4090上运行。同时spatial及fill也通过schnell模型生成，采样步数更少效率更高。 + +所需的都是Comfyui常用的库，一般不需要额外安装 + +有可能缺少的库：diffusers 0.31.0（需要支持FluxPipeline） + +模型读取路径在{comfyui_dir}/models/flux/下 +路径组织结构为： +{comfyui_dir}/models/flux/ + FLUX.1-schnell/ （diffusers格式的flux模型， 下载地址：https://huggingface.co/black-forest-labs/FLUX.1-schnell） + OminiControl/ + depth-anything-small-hf/ （用于depth识别，下载地址：https://huggingface.co/LiheYoung/depth-anything-small-hf/tree/main） + experimental/ （下载地址：https://huggingface.co/Yuanshi/OminiControl/tree/main/experimental） + omini/ （下载地址：https://huggingface.co/Yuanshi/OminiControl/tree/main/omini） \ No newline at end of file diff --git a/rh_config.json b/rh_config.json new file mode 100644 index 0000000..6e0b528 --- /dev/null +++ b/rh_config.json @@ -0,0 +1 @@ +{"enable": true, "untracked_paths": []} diff --git a/rh_omini_fill.py b/rh_omini_fill.py new file mode 100644 index 0000000..4a1f3b0 --- /dev/null +++ b/rh_omini_fill.py @@ -0,0 +1,113 @@ +import torch +from PIL import Image, ImageOps +import numpy as np +from diffusers import FluxPipeline, FluxTransformer2DModel +from ComfyUI_RH_OminiControl.src.generate import generate, seed_everything +from ComfyUI_RH_OminiControl.src.condition import Condition +from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel,T5TokenizerFast +import folder_paths +import os +from ComfyUI_RH_OminiControl.rh_utils import * + +def run(t_img, t_mask, prompt, seed): + + assert t_img.shape[0] == 1 + assert t_mask.shape[0] == 1 + + i = 255. * t_img[0].numpy() + ori_image = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)).convert("RGB").resize((g_width, g_height)) + + mi = 255. * t_mask[0].numpy() + mask = Image.fromarray(np.clip(mi, 0, 255).astype(np.uint8)).convert("L").resize((g_width, g_height)) + mask = ImageOps.invert(mask) + + image = Image.new('RGB', ori_image.size, (0, 0, 0)) + image.paste(ori_image, (0, 0), mask) + + release_gpu() + + flux_dir = os.path.join(folder_paths.models_dir, 'flux', 'FLUX.1-schnell') + lora_model = os.path.join(folder_paths.models_dir, 'flux', 'OminiControl', 'experimental', f'fill.safetensors') + + encoded_condition = encode_condition(flux_dir, image, 'fill') + + text_encoder = CLIPTextModel.from_pretrained( + flux_dir, subfolder="text_encoder", torch_dtype=torch.bfloat16 + ) + text_encoder_2 = T5EncoderModel.from_pretrained( + flux_dir, subfolder="text_encoder_2", torch_dtype=torch.bfloat16 + ) + tokenizer = CLIPTokenizer.from_pretrained(flux_dir, subfolder="tokenizer") + tokenizer_2 = T5TokenizerFast.from_pretrained(flux_dir, subfolder="tokenizer_2") + + pipeline = FluxPipeline.from_pretrained( + flux_dir, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + tokenizer=tokenizer, + tokenizer_2=tokenizer_2, + transformer=None, + vae=None, + ).to("cuda") + + with torch.no_grad(): + prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt( + prompt=prompt, prompt_2=None, max_sequence_length=256 + ) + + del text_encoder + del text_encoder_2 + del tokenizer + del tokenizer_2 + del pipeline + + release_gpu() + + pipeline = FluxPipeline.from_pretrained( + flux_dir, + # transformer=transformer, + text_encoder=None, + text_encoder_2=None, + tokenizer=None, + tokenizer_2=None, + vae=None, + torch_dtype=torch.bfloat16, + ) + + pipeline.to('cuda') + + pipeline.load_lora_weights( + lora_model, + adapter_name='fill', + ) + + condition = Condition('fill', image) + + seed_everything() + + result_latents = generate( + # result_img = generate( + pipeline, + encoded_condition = encoded_condition, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + text_ids=text_ids, + conditions=[condition], + output_type="latent", + return_dict=False, + num_inference_steps=8, + height=g_height, + width=g_width, + ) + + del pipeline + + release_gpu() + + result_img = decode_latents(flux_dir, result_latents[0]).images[0] + + return torch.from_numpy(np.array(result_img).astype(np.float32) / 255.0).unsqueeze(0) + + + + \ No newline at end of file diff --git a/rh_omini_spatial.py b/rh_omini_spatial.py new file mode 100644 index 0000000..3a9c87a --- /dev/null +++ b/rh_omini_spatial.py @@ -0,0 +1,105 @@ +import torch +from PIL import Image +import numpy as np +from diffusers import FluxPipeline, FluxTransformer2DModel +from ComfyUI_RH_OminiControl.src.generate import generate, seed_everything +from ComfyUI_RH_OminiControl.src.condition import Condition +from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel,T5TokenizerFast +import folder_paths +import os +from ComfyUI_RH_OminiControl.rh_utils import * + +def run(t_img, prompt, condition_type, seed): + + assert t_img.shape[0] == 1 + + i = 255. * t_img[0].numpy() + image = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)).convert("RGB").resize((g_width, g_height)) + + release_gpu() + + flux_dir = os.path.join(folder_paths.models_dir, 'flux', 'FLUX.1-schnell') + lora_model = os.path.join(folder_paths.models_dir, 'flux', 'OminiControl', 'experimental', f'{condition_type}.safetensors') + + encoded_condition = encode_condition(flux_dir, image, condition_type) + + text_encoder = CLIPTextModel.from_pretrained( + flux_dir, subfolder="text_encoder", torch_dtype=torch.bfloat16 + ) + text_encoder_2 = T5EncoderModel.from_pretrained( + flux_dir, subfolder="text_encoder_2", torch_dtype=torch.bfloat16 + ) + tokenizer = CLIPTokenizer.from_pretrained(flux_dir, subfolder="tokenizer") + tokenizer_2 = T5TokenizerFast.from_pretrained(flux_dir, subfolder="tokenizer_2") + + pipeline = FluxPipeline.from_pretrained( + flux_dir, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + tokenizer=tokenizer, + tokenizer_2=tokenizer_2, + transformer=None, + vae=None, + ).to("cuda") + + with torch.no_grad(): + prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt( + prompt=prompt, prompt_2=None, max_sequence_length=256 + ) + + del text_encoder + del text_encoder_2 + del tokenizer + del tokenizer_2 + del pipeline + + release_gpu() + + pipeline = FluxPipeline.from_pretrained( + flux_dir, + # transformer=transformer, + text_encoder=None, + text_encoder_2=None, + tokenizer=None, + tokenizer_2=None, + vae=None, + torch_dtype=torch.bfloat16, + ) + + pipeline.to('cuda') + + pipeline.load_lora_weights( + lora_model, + adapter_name=condition_type, + ) + + condition = Condition(condition_type, image) + + seed_everything() + + result_latents = generate( + # result_img = generate( + pipeline, + encoded_condition = encoded_condition, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + text_ids=text_ids, + conditions=[condition], + output_type="latent", + return_dict=False, + num_inference_steps=8, + height=g_height, + width=g_width, + ) + + del pipeline + + release_gpu() + + result_img = decode_latents(flux_dir, result_latents[0]).images[0] + + return torch.from_numpy(np.array(result_img).astype(np.float32) / 255.0).unsqueeze(0) + + + + \ No newline at end of file diff --git a/rh_omini_subject.py b/rh_omini_subject.py new file mode 100644 index 0000000..051322a --- /dev/null +++ b/rh_omini_subject.py @@ -0,0 +1,105 @@ +import torch +from PIL import Image +import numpy as np +from diffusers import FluxPipeline, FluxTransformer2DModel +from ComfyUI_RH_OminiControl.src.generate import generate, seed_everything +from ComfyUI_RH_OminiControl.src.condition import Condition +from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel,T5TokenizerFast +import folder_paths +import os +from ComfyUI_RH_OminiControl.rh_utils import * + +def run(t_img, prompt, seed): + + assert t_img.shape[0] == 1 + + i = 255. * t_img[0].numpy() + image = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)).convert("RGB").resize((g_width, g_height)) + + release_gpu() + + flux_dir = os.path.join(folder_paths.models_dir, 'flux', 'FLUX.1-schnell') + lora_model = os.path.join(folder_paths.models_dir, 'flux', 'OminiControl', 'omini', 'subject_512.safetensors') + + encoded_condition = encode_condition(flux_dir, image) + + text_encoder = CLIPTextModel.from_pretrained( + flux_dir, subfolder="text_encoder", torch_dtype=torch.bfloat16 + ) + text_encoder_2 = T5EncoderModel.from_pretrained( + flux_dir, subfolder="text_encoder_2", torch_dtype=torch.bfloat16 + ) + tokenizer = CLIPTokenizer.from_pretrained(flux_dir, subfolder="tokenizer") + tokenizer_2 = T5TokenizerFast.from_pretrained(flux_dir, subfolder="tokenizer_2") + + pipeline = FluxPipeline.from_pretrained( + flux_dir, + text_encoder=text_encoder, + text_encoder_2=text_encoder_2, + tokenizer=tokenizer, + tokenizer_2=tokenizer_2, + transformer=None, + vae=None, + ).to("cuda") + + with torch.no_grad(): + prompt_embeds, pooled_prompt_embeds, text_ids = pipeline.encode_prompt( + prompt=prompt, prompt_2=None, max_sequence_length=256 + ) + + del text_encoder + del text_encoder_2 + del tokenizer + del tokenizer_2 + del pipeline + + release_gpu() + + pipeline = FluxPipeline.from_pretrained( + flux_dir, + # transformer=transformer, + text_encoder=None, + text_encoder_2=None, + tokenizer=None, + tokenizer_2=None, + vae=None, + torch_dtype=torch.bfloat16, + ) + + pipeline.to('cuda') + + pipeline.load_lora_weights( + lora_model, + adapter_name="subject", + ) + + condition = Condition("subject", image) + + seed_everything() + + result_latents = generate( + # result_img = generate( + pipeline, + encoded_condition = encoded_condition, + prompt_embeds=prompt_embeds, + pooled_prompt_embeds=pooled_prompt_embeds, + text_ids=text_ids, + conditions=[condition], + output_type="latent", + return_dict=False, + num_inference_steps=8, + height=g_height, + width=g_width, + ) + + del pipeline + + release_gpu() + + result_img = decode_latents(flux_dir, result_latents[0]).images[0] + + return torch.from_numpy(np.array(result_img).astype(np.float32) / 255.0).unsqueeze(0) + + + + \ No newline at end of file diff --git a/rh_utils.py b/rh_utils.py new file mode 100644 index 0000000..73e4cd2 --- /dev/null +++ b/rh_utils.py @@ -0,0 +1,58 @@ +import torch +from diffusers import FluxPipeline +import gc +from ComfyUI_RH_OminiControl.src.generate import generate, seed_everything +from ComfyUI_RH_OminiControl.src.condition import Condition + +from diffusers.pipelines.flux.pipeline_flux import ( + FluxPipelineOutput, +) + +g_width = 512 +g_height = 512 + +def release_gpu(): + gc.collect() + torch.cuda.empty_cache() + torch.cuda.reset_max_memory_allocated() + torch.cuda.reset_peak_memory_stats() + +def encode_condition(flux_dir, image, condition_type='subject'): + pipeline = FluxPipeline.from_pretrained( + flux_dir, + text_encoder=None, + text_encoder_2=None, + tokenizer=None, + tokenizer_2=None, + transformer=None, + torch_dtype=torch.bfloat16, + ).to("cuda") + + condition = Condition(condition_type, image) + tokens, ids, type_id = condition.encode(pipeline) + + del condition + del pipeline + release_gpu() + + return (tokens, ids, type_id) + +def decode_latents(flux_dir, latents): + pipeline = FluxPipeline.from_pretrained( + flux_dir, + text_encoder=None, + text_encoder_2=None, + tokenizer=None, + tokenizer_2=None, + transformer=None, + torch_dtype=torch.bfloat16, + ).to("cuda") + + latents = pipeline._unpack_latents(latents, g_width, g_height, pipeline.vae_scale_factor) + latents = ( + latents / pipeline.vae.config.scaling_factor + ) + pipeline.vae.config.shift_factor + image = pipeline.vae.decode(latents, return_dict=False)[0] + image = pipeline.image_processor.postprocess(image, output_type="pil") + + return FluxPipelineOutput(images=image) \ No newline at end of file