mirror of
https://github.com/lllyasviel/Fooocus.git
synced 2026-08-16 13:13:16 +02:00
2.0.80 (#520)
* Rework many patches and some UI details. * Speed up processing. * Move Colab to independent branch. * Implemented CFG Scale and TSNR correction when CFG is bigger than 10. * Implemented Developer Mode with more options to debug.
This commit is contained in:
lllyasviel
+68
-50
@@ -10,13 +10,15 @@ import torch
|
||||
import numpy as np
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.model_detection
|
||||
import comfy.model_patcher
|
||||
import comfy.utils
|
||||
|
||||
from comfy.sd import load_checkpoint_guess_config
|
||||
from nodes import VAEDecode, EmptyLatentImage, VAEEncode, VAEEncodeTiled, VAEDecodeTiled, VAEEncodeForInpaint
|
||||
from comfy.sample import prepare_mask, broadcast_cond, load_additional_models, cleanup_additional_models
|
||||
from comfy.model_base import SDXLRefiner
|
||||
from comfy.sd import model_lora_keys_unet, model_lora_keys_clip, load_lora
|
||||
from comfy.sample import prepare_mask, broadcast_cond, get_additional_models, cleanup_additional_models
|
||||
from modules.patch import patched_sampler_cfg_function, patched_model_function_wrapper
|
||||
from comfy.lora import model_lora_keys_unet, model_lora_keys_clip, load_lora
|
||||
from modules.samplers_advanced import KSamplerBasic, KSamplerWithRefiner
|
||||
|
||||
|
||||
@@ -29,34 +31,61 @@ opVAEEncodeForInpaint = VAEEncodeForInpaint()
|
||||
|
||||
|
||||
class StableDiffusionModel:
|
||||
def __init__(self, unet, vae, clip, clip_vision, model_filename=None):
|
||||
if isinstance(model_filename, str):
|
||||
is_refiner = isinstance(unet.model, SDXLRefiner)
|
||||
if unet is not None:
|
||||
unet.model.model_file = dict(filename=model_filename, prefix='model')
|
||||
if clip is not None:
|
||||
clip.cond_stage_model.model_file = dict(filename=model_filename, prefix='refiner_clip' if is_refiner else 'base_clip')
|
||||
if vae is not None:
|
||||
vae.first_stage_model.model_file = dict(filename=model_filename, prefix='first_stage_model')
|
||||
def __init__(self, unet, vae, clip, clip_vision):
|
||||
self.unet = unet
|
||||
self.vae = vae
|
||||
self.clip = clip
|
||||
self.clip_vision = clip_vision
|
||||
|
||||
def to_meta(self):
|
||||
if self.unet is not None:
|
||||
self.unet.model.to('meta')
|
||||
if self.clip is not None:
|
||||
self.clip.cond_stage_model.to('meta')
|
||||
if self.vae is not None:
|
||||
self.vae.first_stage_model.to('meta')
|
||||
|
||||
@torch.no_grad()
|
||||
@torch.inference_mode()
|
||||
def load_unet_only(unet_path):
|
||||
sd_raw = comfy.utils.load_torch_file(unet_path)
|
||||
sd = {}
|
||||
flag = 'model.diffusion_model.'
|
||||
for k in list(sd_raw.keys()):
|
||||
if k.startswith(flag):
|
||||
sd[k[len(flag):]] = sd_raw[k]
|
||||
del sd_raw[k]
|
||||
|
||||
parameters = comfy.utils.calculate_parameters(sd)
|
||||
fp16 = comfy.model_management.should_use_fp16(model_params=parameters)
|
||||
if "input_blocks.0.0.weight" in sd:
|
||||
# ldm
|
||||
model_config = comfy.model_detection.model_config_from_unet(sd, "", fp16)
|
||||
if model_config is None:
|
||||
raise RuntimeError("ERROR: Could not detect model type of: {}".format(unet_path))
|
||||
new_sd = sd
|
||||
else:
|
||||
# diffusers
|
||||
model_config = comfy.model_detection.model_config_from_diffusers_unet(sd, fp16)
|
||||
if model_config is None:
|
||||
print("ERROR UNSUPPORTED UNET", unet_path)
|
||||
return None
|
||||
|
||||
diffusers_keys = comfy.utils.unet_to_diffusers(model_config.unet_config)
|
||||
|
||||
new_sd = {}
|
||||
for k in diffusers_keys:
|
||||
if k in sd:
|
||||
new_sd[diffusers_keys[k]] = sd.pop(k)
|
||||
else:
|
||||
print(diffusers_keys[k], k)
|
||||
offload_device = comfy.model_management.unet_offload_device()
|
||||
model = model_config.get_model(new_sd, "")
|
||||
model = model.to(offload_device)
|
||||
model.load_model_weights(new_sd, "")
|
||||
return comfy.model_patcher.ModelPatcher(model, load_device=comfy.model_management.get_torch_device(), offload_device=offload_device)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
@torch.inference_mode()
|
||||
def load_model(ckpt_filename):
|
||||
unet, clip, vae, clip_vision = load_checkpoint_guess_config(ckpt_filename)
|
||||
return StableDiffusionModel(unet=unet, clip=clip, vae=vae, clip_vision=clip_vision, model_filename=ckpt_filename)
|
||||
unet.model_options['sampler_cfg_function'] = patched_sampler_cfg_function
|
||||
unet.model_options['model_function_wrapper'] = patched_model_function_wrapper
|
||||
return StableDiffusionModel(unet=unet, clip=clip, vae=vae, clip_vision=clip_vision)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
@@ -74,20 +103,19 @@ def load_sd_lora(model, lora_filename, strength_model=1.0, strength_clip=1.0):
|
||||
key_map = model_lora_keys_clip(model.clip.cond_stage_model, key_map)
|
||||
loaded = load_lora(lora, key_map)
|
||||
|
||||
new_modelpatcher = model.unet.clone()
|
||||
k = new_modelpatcher.add_patches(loaded, strength_model)
|
||||
new_unet = model.unet.clone()
|
||||
loaded_unet_keys = new_unet.add_patches(loaded, strength_model)
|
||||
|
||||
new_clip = model.clip.clone()
|
||||
k1 = new_clip.add_patches(loaded, strength_clip)
|
||||
loaded_clip_keys = new_clip.add_patches(loaded, strength_clip)
|
||||
|
||||
loaded_keys = set(list(loaded_unet_keys) + list(loaded_clip_keys))
|
||||
|
||||
k = set(k)
|
||||
k1 = set(k1)
|
||||
for x in loaded:
|
||||
if (x not in k) and (x not in k1):
|
||||
print("Lora missed: ", x)
|
||||
if x not in loaded_keys:
|
||||
print("Lora key not loaded: ", x)
|
||||
|
||||
unet, clip = new_modelpatcher, new_clip
|
||||
return StableDiffusionModel(unet=unet, clip=clip, vae=model.vae, clip_vision=model.clip_vision)
|
||||
return StableDiffusionModel(unet=new_unet, clip=new_clip, vae=model.vae, clip_vision=model.clip_vision)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
@@ -142,7 +170,7 @@ VAE_approx_model = None
|
||||
|
||||
@torch.no_grad()
|
||||
@torch.inference_mode()
|
||||
def get_previewer(device, latent_format):
|
||||
def get_previewer():
|
||||
global VAE_approx_model
|
||||
|
||||
if VAE_approx_model is None:
|
||||
@@ -181,12 +209,7 @@ def get_previewer(device, latent_format):
|
||||
def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sampler_name='dpmpp_fooocus_2m_sde_inpaint_seamless',
|
||||
scheduler='karras', denoise=1.0, disable_noise=False, start_step=None, last_step=None,
|
||||
force_full_denoise=False, callback_function=None):
|
||||
# SCHEDULERS = ["normal", "karras", "exponential", "simple", "ddim_uniform"]
|
||||
# SAMPLERS = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral",
|
||||
# "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
|
||||
# "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "ddim", "uni_pc", "uni_pc_bh2"]
|
||||
|
||||
seed = seed if isinstance(seed, int) else random.randint(1, 2 ** 64)
|
||||
seed = seed if isinstance(seed, int) else random.randint(0, 2**63 - 1)
|
||||
|
||||
device = comfy.model_management.get_torch_device()
|
||||
latent_image = latent["samples"]
|
||||
@@ -201,11 +224,12 @@ def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sa
|
||||
if "noise_mask" in latent:
|
||||
noise_mask = latent["noise_mask"]
|
||||
|
||||
previewer = get_previewer(device, model.model.latent_format)
|
||||
previewer = get_previewer()
|
||||
|
||||
pbar = comfy.utils.ProgressBar(steps)
|
||||
|
||||
def callback(step, x0, x, total_steps):
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
y = None
|
||||
if previewer is not None:
|
||||
y = previewer(x0, step, total_steps)
|
||||
@@ -219,7 +243,8 @@ def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sa
|
||||
if noise_mask is not None:
|
||||
noise_mask = prepare_mask(noise_mask, noise.shape, device)
|
||||
|
||||
comfy.model_management.load_model_gpu(model)
|
||||
models, inference_memory = get_additional_models(positive, negative, model.model_dtype())
|
||||
comfy.model_management.load_models_gpu([model] + models, comfy.model_management.batch_area_memory(noise.shape[0] * noise.shape[2] * noise.shape[3]) + inference_memory)
|
||||
real_model = model.model
|
||||
|
||||
noise = noise.to(device)
|
||||
@@ -228,8 +253,6 @@ def ksampler(model, positive, negative, latent, seed=None, steps=30, cfg=7.0, sa
|
||||
positive_copy = broadcast_cond(positive, noise.shape[0], device)
|
||||
negative_copy = broadcast_cond(negative, noise.shape[0], device)
|
||||
|
||||
models = load_additional_models(positive, negative, model.model_dtype())
|
||||
|
||||
sampler = KSamplerBasic(real_model, steps=steps, device=device, sampler=sampler_name, scheduler=scheduler,
|
||||
denoise=denoise, model_options=model.model_options)
|
||||
|
||||
@@ -254,12 +277,7 @@ def ksampler_with_refiner(model, positive, negative, refiner, refiner_positive,
|
||||
seed=None, steps=30, refiner_switch_step=20, cfg=7.0, sampler_name='dpmpp_fooocus_2m_sde_inpaint_seamless',
|
||||
scheduler='karras', denoise=1.0, disable_noise=False, start_step=None, last_step=None,
|
||||
force_full_denoise=False, callback_function=None):
|
||||
# SCHEDULERS = ["normal", "karras", "exponential", "simple", "ddim_uniform"]
|
||||
# SAMPLERS = ["euler", "euler_ancestral", "heun", "dpm_2", "dpm_2_ancestral",
|
||||
# "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
|
||||
# "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "ddim", "uni_pc", "uni_pc_bh2"]
|
||||
|
||||
seed = seed if isinstance(seed, int) else random.randint(1, 2 ** 64)
|
||||
seed = seed if isinstance(seed, int) else random.randint(0, 2**63 - 1)
|
||||
|
||||
device = comfy.model_management.get_torch_device()
|
||||
latent_image = latent["samples"]
|
||||
@@ -274,11 +292,12 @@ def ksampler_with_refiner(model, positive, negative, refiner, refiner_positive,
|
||||
if "noise_mask" in latent:
|
||||
noise_mask = latent["noise_mask"]
|
||||
|
||||
previewer = get_previewer(device, model.model.latent_format)
|
||||
previewer = get_previewer()
|
||||
|
||||
pbar = comfy.utils.ProgressBar(steps)
|
||||
|
||||
def callback(step, x0, x, total_steps):
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
y = None
|
||||
if previewer is not None:
|
||||
y = previewer(x0, step, total_steps)
|
||||
@@ -292,7 +311,8 @@ def ksampler_with_refiner(model, positive, negative, refiner, refiner_positive,
|
||||
if noise_mask is not None:
|
||||
noise_mask = prepare_mask(noise_mask, noise.shape, device)
|
||||
|
||||
comfy.model_management.load_model_gpu(model)
|
||||
models, inference_memory = get_additional_models(positive, negative, model.model_dtype())
|
||||
comfy.model_management.load_models_gpu([model] + models, comfy.model_management.batch_area_memory(noise.shape[0] * noise.shape[2] * noise.shape[3]) + inference_memory)
|
||||
|
||||
noise = noise.to(device)
|
||||
latent_image = latent_image.to(device)
|
||||
@@ -303,8 +323,6 @@ def ksampler_with_refiner(model, positive, negative, refiner, refiner_positive,
|
||||
refiner_positive_copy = broadcast_cond(refiner_positive, noise.shape[0], device)
|
||||
refiner_negative_copy = broadcast_cond(refiner_negative, noise.shape[0], device)
|
||||
|
||||
models = load_additional_models(positive, negative, model.model_dtype())
|
||||
|
||||
sampler = KSamplerWithRefiner(model=model, refiner_model=refiner, steps=steps, device=device,
|
||||
sampler=sampler_name, scheduler=scheduler,
|
||||
denoise=denoise, model_options=model.model_options)
|
||||
|
||||
Reference in New Issue
Block a user