|
import sys |
|
import os |
|
import torch |
|
|
|
from PIL import Image |
|
from typing import List |
|
import numpy as np |
|
from utils import ( |
|
tensor_to_pil, |
|
pil_to_tensor, |
|
pad_image, |
|
postprocess_image, |
|
preprocess_image, |
|
downloadModels, |
|
) |
|
|
|
sys.path.append(os.path.dirname("./ComfyUI/")) |
|
from ComfyUI.nodes import ( |
|
CheckpointLoaderSimple, |
|
VAEDecode, |
|
VAEEncode, |
|
KSampler, |
|
EmptyLatentImage, |
|
CLIPTextEncode, |
|
) |
|
from ComfyUI.comfy_extras.nodes_compositing import JoinImageWithAlpha |
|
from ComfyUI.comfy_extras.nodes_mask import InvertMask, MaskToImage |
|
|
|
from ComfyUI.comfy import samplers |
|
|
|
from ComfyUI.custom_nodes.layerdiffuse.layered_diffusion import ( |
|
LayeredDiffusionFG, |
|
LayeredDiffusionDecode, |
|
LayeredDiffusionCond, |
|
) |
|
import gradio as gr |
|
from briarmbg import BriaRMBG |
|
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
|
downloadModels() |
|
|
|
with torch.inference_mode(): |
|
ckpt_load_checkpoint = CheckpointLoaderSimple().load_checkpoint |
|
ckpt = ckpt_load_checkpoint(ckpt_name="juggernautXL_v8Rundiffusion.safetensors") |
|
|
|
cliptextencode = CLIPTextEncode().encode |
|
emptylatentimage_generate = EmptyLatentImage().generate |
|
ksampler_sample = KSampler().sample |
|
vae_decode = VAEDecode().decode |
|
vae_encode = VAEEncode().encode |
|
ld_fg_apply_layered_diffusion = LayeredDiffusionFG().apply_layered_diffusion |
|
ld_cond_apply_layered_diffusion = LayeredDiffusionCond().apply_layered_diffusion |
|
|
|
ld_decode = LayeredDiffusionDecode().decode |
|
mask_to_image = MaskToImage().mask_to_image |
|
invert_mask = InvertMask().invert |
|
join_image_with_alpha = JoinImageWithAlpha().join_image_with_alpha |
|
rmbg_model = BriaRMBG.from_pretrained("briaai/RMBG-1.4").to(device) |
|
|
|
|
|
def predict( |
|
prompt: str, |
|
negative_prompt: str, |
|
input_image: Image.Image | None, |
|
remove_bg: bool, |
|
cond_mode: str, |
|
seed: int, |
|
sampler_name: str, |
|
scheduler: str, |
|
steps: int, |
|
cfg: float, |
|
denoise: float, |
|
): |
|
try: |
|
with torch.inference_mode(): |
|
cliptextencode_prompt = cliptextencode( |
|
text=prompt, |
|
clip=ckpt[1], |
|
) |
|
cliptextencode_negative_prompt = cliptextencode( |
|
text=negative_prompt, |
|
clip=ckpt[1], |
|
) |
|
emptylatentimage_sample = emptylatentimage_generate( |
|
width=1024, height=1024, batch_size=1 |
|
) |
|
|
|
if input_image is not None: |
|
input_image = pad_image(input_image).resize((1024, 1024)) |
|
if remove_bg: |
|
orig_im_size = input_image.size |
|
image = preprocess_image(np.array(input_image), [1024, 1024]).to( |
|
device |
|
) |
|
|
|
result = rmbg_model(image) |
|
|
|
result_mask_image = postprocess_image(result[0][0], orig_im_size) |
|
|
|
|
|
pil_mask = Image.fromarray(result_mask_image) |
|
no_bg_image = Image.new("RGBA", pil_mask.size, (0, 0, 0, 0)) |
|
no_bg_image.paste(input_image, mask=pil_mask) |
|
input_image = no_bg_image |
|
|
|
img_tensor = pil_to_tensor(input_image) |
|
img_latent = vae_encode(pixels=img_tensor[0], vae=ckpt[2]) |
|
layereddiffusionapply_sample = ld_cond_apply_layered_diffusion( |
|
config=cond_mode, |
|
weight=1, |
|
model=ckpt[0], |
|
cond=cliptextencode_prompt[0], |
|
uncond=cliptextencode_negative_prompt[0], |
|
latent=img_latent[0], |
|
) |
|
ksampler = ksampler_sample( |
|
steps=steps, |
|
cfg=cfg, |
|
sampler_name=sampler_name, |
|
scheduler=scheduler, |
|
seed=seed, |
|
model=layereddiffusionapply_sample[0], |
|
positive=layereddiffusionapply_sample[1], |
|
negative=layereddiffusionapply_sample[2], |
|
latent_image=emptylatentimage_sample[0], |
|
denoise=denoise, |
|
) |
|
|
|
vaedecode_sample = vae_decode( |
|
samples=ksampler[0], |
|
vae=ckpt[2], |
|
) |
|
layereddiffusiondecode_sample = ld_decode( |
|
sd_version="SDXL", |
|
sub_batch_size=16, |
|
samples=ksampler[0], |
|
images=vaedecode_sample[0], |
|
) |
|
|
|
rgb_img = tensor_to_pil(vaedecode_sample[0]) |
|
return flatten([rgb_img]) |
|
else: |
|
layereddiffusionapply_sample = ld_fg_apply_layered_diffusion( |
|
config="SDXL, Conv Injection", weight=1, model=ckpt[0] |
|
) |
|
ksampler = ksampler_sample( |
|
steps=steps, |
|
cfg=cfg, |
|
sampler_name=sampler_name, |
|
scheduler=scheduler, |
|
seed=seed, |
|
model=layereddiffusionapply_sample[0], |
|
positive=cliptextencode_prompt[0], |
|
negative=cliptextencode_negative_prompt[0], |
|
latent_image=emptylatentimage_sample[0], |
|
denoise=denoise, |
|
) |
|
|
|
vaedecode_sample = vae_decode( |
|
samples=ksampler[0], |
|
vae=ckpt[2], |
|
) |
|
layereddiffusiondecode_sample = ld_decode( |
|
sd_version="SDXL", |
|
sub_batch_size=16, |
|
samples=ksampler[0], |
|
images=vaedecode_sample[0], |
|
) |
|
mask = mask_to_image(mask=layereddiffusiondecode_sample[1]) |
|
ld_image = tensor_to_pil(layereddiffusiondecode_sample[0][0]) |
|
inverted_mask = invert_mask(mask=layereddiffusiondecode_sample[1]) |
|
rgba_img = join_image_with_alpha( |
|
image=layereddiffusiondecode_sample[0], alpha=inverted_mask[0] |
|
) |
|
rgba_img = tensor_to_pil(rgba_img[0]) |
|
mask = tensor_to_pil(mask[0]) |
|
rgb_img = tensor_to_pil(vaedecode_sample[0]) |
|
|
|
return flatten([rgba_img, mask]) |
|
|
|
except Exception as e: |
|
raise gr.Error(e) |
|
|
|
|
|
examples = [["An old men sit on a chair looking at the sky"]] |
|
|
|
|
|
def flatten(l: List[List[any]]) -> List[any]: |
|
return [item for sublist in l for item in sublist] |
|
|
|
|
|
def predict_examples(prompt, negative_prompt): |
|
return predict( |
|
prompt, negative_prompt, None, False, None, 0, "euler", "normal", 20, 8.0, 1.0 |
|
) |
|
|
|
|
|
css = """ |
|
.gradio-container{ |
|
max-width: 50rem; |
|
} |
|
""" |
|
with gr.Blocks(css=css) as blocks: |
|
gr.Markdown("""# LayerDiffuse (unofficial) |
|
Using ComfyUI building blocks with custom node by [huchenlei](https://github.com/huchenlei/ComfyUI-layerdiffuse) |
|
""") |
|
|
|
with gr.Row(): |
|
with gr.Column(): |
|
prompt = gr.Text(label="Prompt") |
|
negative_prompt = gr.Text(label="Negative Prompt") |
|
button = gr.Button("Generate") |
|
with gr.Accordion(open=False, label="Input Images (Optional)"): |
|
with gr.Group(): |
|
cond_mode = gr.Radio( |
|
value="SDXL, Foreground", |
|
choices=["SDXL, Foreground", "SDXL, Background"], |
|
info="Whether to use input image as foreground or background", |
|
) |
|
remove_bg = gr.Checkbox( |
|
info="Remove background using BriaRMBG", |
|
label="Remove Background", |
|
value=False, |
|
) |
|
input_image = gr.Image(label="Input Image", type="pil") |
|
with gr.Accordion(open=False, label="Advanced Options"): |
|
seed = gr.Slider( |
|
label="Seed", |
|
value=0, |
|
minimum=-1, |
|
maximum=0xFFFFFFFFFFFFFFFF, |
|
step=1, |
|
randomize=True, |
|
) |
|
sampler_name = gr.Dropdown( |
|
choices=samplers.KSampler.SAMPLERS, |
|
label="Sampler Name", |
|
value=samplers.KSampler.SAMPLERS[0], |
|
) |
|
scheduler = gr.Dropdown( |
|
choices=samplers.KSampler.SCHEDULERS, |
|
label="Scheduler", |
|
value=samplers.KSampler.SCHEDULERS[0], |
|
) |
|
steps = gr.Slider( |
|
label="Steps", value=20, minimum=1, maximum=30, step=1 |
|
) |
|
cfg = gr.Number( |
|
label="CFG", value=8.0, minimum=0.0, maximum=100.0, step=0.1 |
|
) |
|
denoise = gr.Number( |
|
label="Denoise", value=1.0, minimum=0.0, maximum=1.0, step=0.01 |
|
) |
|
|
|
with gr.Column(scale=1.8): |
|
gallery = gr.Gallery(columns=[2], object_fit="contain", height="unset") |
|
|
|
inputs = [ |
|
prompt, |
|
negative_prompt, |
|
input_image, |
|
remove_bg, |
|
cond_mode, |
|
seed, |
|
sampler_name, |
|
scheduler, |
|
steps, |
|
cfg, |
|
denoise, |
|
] |
|
outputs = [gallery] |
|
|
|
gr.Examples( |
|
fn=predict_examples, |
|
examples=examples, |
|
inputs=[prompt, negative_prompt], |
|
outputs=outputs, |
|
cache_examples=False, |
|
) |
|
|
|
button.click(fn=predict, inputs=inputs, outputs=outputs) |
|
|
|
|
|
if __name__ == "__main__": |
|
blocks.launch() |
|
|
