ComfyUI-DiffusersImageOutpaint/utils.py at main · GiusTex/ComfyUI-DiffusersImageOutpaint · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
import torch
import gc
import os
import numpy as np
import json
import comfy.model_management as mm
from PIL import Image

from folder_paths import map_legacy, folder_names_and_paths
from .pipeline_fill_sd_xl import StableDiffusionXLFillPipeline


def get_first_folder_list(folder_name: str) -> tuple[list[str], dict[str, float], float]:
    folder_name = map_legacy(folder_name)
    global folder_names_and_paths
    folders = folder_names_and_paths[folder_name]
    if folder_name == "unet":
        root_folder = folders[0][0]
    elif folder_name == "diffusion_models":
        root_folder = folders[0][1]
    elif folder_name == "controlnet":
        root_folder = folders[0][0]
    visible_folders = [name for name in os.listdir(root_folder) if os.path.isdir(os.path.join(root_folder, name))]
    return visible_folders

def get_config_folder_list(folder_name: str) -> tuple[list[str], dict[str, float], float]:
    my_dir = os.path.dirname(os.path.abspath(__file__))
    configs_dir = f"{my_dir}/{folder_name}"

    folders = [f for f in os.listdir(configs_dir) if os.path.isdir(os.path.join(configs_dir, f))]
    return folders

# Tensor to PIL (grabbed from WAS Suite)
def tensor2pil(image: torch.Tensor) -> Image.Image:
    return Image.fromarray(np.clip(255. * image.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))

# Convert PIL to Tensor (grabbed from WAS Suite)
def pil2tensor(image: Image.Image) -> torch.Tensor:
    return torch.from_numpy(np.array(image).astype(np.float32) / 255.0).unsqueeze(0)


def get_device_by_name(device):
    if device == 'auto':
        device = mm.get_torch_device()
    return device


def get_dtype_by_name(dtype):
    if dtype == 'auto':
        if mm.should_use_fp16():
            dtype = torch.float16
        elif mm.should_use_bf16():
            dtype = torch.bfloat16
        else:
            dtype = torch.float32
    elif dtype== "fp16":
        dtype = torch.float16
    elif dtype == "bf16":
        dtype = torch.bfloat16
    elif dtype == "fp32":
        dtype = torch.float32
    elif dtype == "fp8_e4m3fn":
        dtype = torch.float8_e4m3fn
    elif dtype == "fp8_e4m3fnuz":
        dtype = torch.float8_e4m3fnuz
    elif dtype == "fp8_e5m2":
        dtype = torch.float8_e5m2
    elif dtype == "fp8_e5m2fnuz":
        dtype = torch.float8_e5m2fnuz

    return dtype

def clearVram(device):
    gc.collect()

    if device.type == "cuda":
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()
    elif device.type == "mps":
        torch.mps.empty_cache()
    elif device.type == "xla":
        torch.xla.empty_cache()
    elif device.type == "xpu":
        torch.xpu.empty_cache()
    elif device.type == "meta":
        torch.meta.empty_cache()


class TCDScheduler_Custom:
    def __init__(self, **kwargs):
        for key, value in kwargs.items():
            setattr(self, key, value)

    def scale_model_input(self, input, t):
        scale_factor = getattr(self, 'scale_factor', 1)
        return input * scale_factor

    def __repr__(self):
        attrs = {key: value for key, value in self.__dict__.items()}
        return f"TCDScheduler({attrs})"


def test_scheduler_scale_model_input(comfy_dir, model_type):
    scheduler_config_path = f"{comfy_dir}/custom_nodes/ComfyUI-DiffusersImageOutpaint/configs/{model_type}/scheduler/scheduler_config.json"

    with open(scheduler_config_path, 'r') as f:
        config = json.load(f)

    scheduler = TCDScheduler_Custom(**config)
    scale_model_input_method = scheduler.scale_model_input

    return scale_model_input_method


def diffuserOutpaintSamples(device, dtype, keep_model_device, scheduler, scale_model_input_method, model, control_net, positive, negative,
                                                      cnet_image, controlnet_strength, guidance_scale, steps):

    prompt_embeds = positive["prompt_embeds"]
    pooled_prompt_embeds = positive["pooled_prompt_embeds"]
    negative_prompt_embeds = negative["prompt_embeds"]
    negative_pooled_prompt_embeds = negative["pooled_prompt_embeds"]
    controlnet_model = control_net

    device = get_device_by_name(device)
    dtype = get_dtype_by_name(dtype)

    timesteps = None
    unet = model

    pipe = StableDiffusionXLFillPipeline()

    rgb_latents = list(pipe(
            prompt_embeds=prompt_embeds,
            negative_prompt_embeds=negative_prompt_embeds,
            pooled_prompt_embeds=pooled_prompt_embeds,
            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
            image=cnet_image,
            num_inference_steps=steps,
            controlnet_model=controlnet_model,
            controlnet_conditioning_scale=controlnet_strength,
            guidance_scale=guidance_scale,
            device=device,
            dtype=dtype,
            unet=unet,
            timesteps=timesteps,
            scale_model_input_method=scale_model_input_method,
            keep_model_device=keep_model_device,
            scheduler=scheduler,
    ))

    last_rgb_latent = rgb_latents[-1] # Access the last image

    del pipe, unet, controlnet_model, scheduler, prompt_embeds, negative_prompt_embeds, pooled_prompt_embeds, negative_pooled_prompt_embeds

    clearVram(device)

    return last_rgb_latent