init
This commit is contained in:
HM-RunningHub
@@ -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']
|
||||
@@ -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,
|
||||
}
|
||||
@@ -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）
|
||||
@@ -0,0 +1 @@
|
||||
{"enable": true, "untracked_paths": []}
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
|
||||
+58
@@ -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)
|
||||
Reference in New Issue
Block a user