app.py

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,
    examples,
)

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_version6Rundiffusion.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,
    remove_bg: bool,
    cond_mode: str,
    seed: int,
    sampler_name: str,
    scheduler: str,
    steps: int,
    cfg: float,
    denoise: float,
):
    seed = seed if seed != -1 else np.random.randint(0, 2**63 - 1)
    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)
                    # post process
                    result_mask_image = postprocess_image(result[0][0], orig_im_size)

                    # save result
                    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 (rgb_img[0], rgb_img[0], seed)

            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 (rgba_img[0], mask[0], seed)
                # return flatten([rgba_img, mask, rgb_img, ld_image])
    except Exception as e:
        raise gr.Error(e)


def flatten(l: List[List[any]]) -> List[any]:
    return [item for sublist in l for item in sublist]


def predict_examples(
    prompt,
    negative_prompt,
    input_image=None,
    remove_bg=False,
    cond_mode=None,
    seed=-1,
    cfg=10,
):
    return predict(
        prompt,
        negative_prompt,
        input_image,
        remove_bg,
        cond_mode,
        seed,
        "dpmpp_2m_sde_gpu",
        "karras",
        30,
        cfg,
        1.0,
    )


css = """
.gradio-container{
    max-width: 85rem !important;
}
"""
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)  
    models: [LayerDiffusion/layerdiffusion-v1](https://huggingface.co/LayerDiffusion/layerdiffusion-v1/tree/main)
                """)

    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"):
                with gr.Group():
                    with gr.Row():
                        seed = gr.Slider(
                            label="Seed",
                            value=-1,
                            minimum=-1,
                            maximum=0xFFFFFFFFFFFFFFFF,
                            step=1,
                        )
                        curr_seed = gr.Number(
                            value=-1, interactive=False, scale=0, label=" "
                        )
                sampler_name = gr.Dropdown(
                    choices=samplers.KSampler.SAMPLERS,
                    label="Sampler Name",
                    value="dpmpp_2m_sde_gpu",
                )
                scheduler = gr.Dropdown(
                    choices=samplers.KSampler.SCHEDULERS,
                    label="Scheduler",
                    value="karras",
                )
                steps = gr.Slider(
                    label="Steps", value=20, minimum=1, maximum=50, step=1
                )
                cfg = gr.Number(
                    label="CFG", value=5.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():
            image = gr.Image()
            with gr.Accordion(label="Mask", open=False):
                mask = gr.Image()

    inputs = [
        prompt,
        negative_prompt,
        input_image,
        remove_bg,
        cond_mode,
        seed,
        sampler_name,
        scheduler,
        steps,
        cfg,
        denoise,
    ]
    outputs = [image, mask, curr_seed]

    gr.Examples(
        fn=predict_examples,
        examples=examples,
        inputs=[prompt, negative_prompt, input_image, remove_bg, cond_mode, seed],
        outputs=outputs,
        cache_examples=True,
    )

    button.click(fn=predict, inputs=inputs, outputs=outputs)


if __name__ == "__main__":
    blocks.launch()