mirror of
https://github.com/lllyasviel/Fooocus.git
synced 2026-08-16 13:13:16 +02:00
rework refiner for some potential new features (#642)
* sync * sync * sync * sync * sync * sync * sync * sync * sync * sync * sync * sync * sync * sync * sync * sync * sync * sync * sync * sync * sync * sync * sync * sync
This commit is contained in:
lllyasviel
+161
-29
@@ -4,9 +4,9 @@ import torch
|
||||
import modules.path
|
||||
import comfy.model_management
|
||||
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
from comfy.model_base import SDXL, SDXLRefiner
|
||||
from modules.expansion import FooocusExpansion
|
||||
from modules.sample_hijack import clip_separate
|
||||
|
||||
|
||||
xl_base: core.StableDiffusionModel = None
|
||||
@@ -15,14 +15,15 @@ xl_base_hash = ''
|
||||
xl_base_patched: core.StableDiffusionModel = None
|
||||
xl_base_patched_hash = ''
|
||||
|
||||
xl_refiner: ModelPatcher = None
|
||||
xl_refiner: core.StableDiffusionModel = None
|
||||
xl_refiner_hash = ''
|
||||
|
||||
final_expansion = None
|
||||
final_unet = None
|
||||
final_clip = None
|
||||
final_vae = None
|
||||
final_refiner = None
|
||||
final_refiner_unet = None
|
||||
final_refiner_vae = None
|
||||
|
||||
loaded_ControlNets = {}
|
||||
|
||||
@@ -60,8 +61,10 @@ def assert_model_integrity():
|
||||
error_message = 'You have selected base model other than SDXL. This is not supported yet.'
|
||||
|
||||
if xl_refiner is not None:
|
||||
if not isinstance(xl_refiner.model, SDXLRefiner):
|
||||
error_message = 'You have selected refiner model other than SDXL refiner. This is not supported yet.'
|
||||
if xl_refiner.unet is None or xl_refiner.unet.model is None:
|
||||
error_message = 'You have selected an invalid refiner!'
|
||||
elif not isinstance(xl_refiner.unet.model, SDXL) and not isinstance(xl_refiner.unet.model, SDXLRefiner):
|
||||
error_message = 'SD1.5 or 2.1 as refiner is not supported!'
|
||||
|
||||
if error_message is not None:
|
||||
raise NotImplementedError(error_message)
|
||||
@@ -109,9 +112,19 @@ def refresh_refiner_model(name):
|
||||
print(f'Refiner unloaded.')
|
||||
return
|
||||
|
||||
xl_refiner = core.load_unet_only(filename)
|
||||
xl_refiner = core.load_model(filename)
|
||||
xl_refiner_hash = model_hash
|
||||
print(f'Refiner model loaded: {model_hash}')
|
||||
|
||||
if isinstance(xl_refiner.unet.model, SDXL):
|
||||
xl_refiner.clip = None
|
||||
xl_refiner.vae = None
|
||||
elif isinstance(xl_refiner.unet.model, SDXLRefiner):
|
||||
xl_refiner.clip = None
|
||||
xl_refiner.vae = None
|
||||
else:
|
||||
xl_refiner = None # 1.5/2.1 not supported yet.
|
||||
|
||||
return
|
||||
|
||||
|
||||
@@ -203,15 +216,23 @@ def prepare_text_encoder(async_call=True):
|
||||
@torch.no_grad()
|
||||
@torch.inference_mode()
|
||||
def refresh_everything(refiner_model_name, base_model_name, loras):
|
||||
global final_unet, final_clip, final_vae, final_refiner, final_expansion
|
||||
global final_unet, final_clip, final_vae, final_refiner_unet, final_refiner_vae, final_expansion
|
||||
|
||||
refresh_refiner_model(refiner_model_name)
|
||||
refresh_base_model(base_model_name)
|
||||
refresh_loras(loras)
|
||||
assert_model_integrity()
|
||||
|
||||
final_unet, final_clip, final_vae, final_refiner = \
|
||||
xl_base_patched.unet, xl_base_patched.clip, xl_base_patched.vae, xl_refiner
|
||||
final_unet = xl_base_patched.unet
|
||||
final_clip = xl_base_patched.clip
|
||||
final_vae = xl_base_patched.vae
|
||||
|
||||
if xl_refiner is None:
|
||||
final_refiner_unet = None
|
||||
final_refiner_vae = None
|
||||
else:
|
||||
final_refiner_unet = xl_refiner.unet
|
||||
final_refiner_vae = xl_refiner.vae
|
||||
|
||||
if final_expansion is None:
|
||||
final_expansion = FooocusExpansion()
|
||||
@@ -236,30 +257,141 @@ refresh_everything(
|
||||
|
||||
@torch.no_grad()
|
||||
@torch.inference_mode()
|
||||
def process_diffusion(positive_cond, negative_cond, steps, switch, width, height, image_seed, callback, sampler_name, scheduler_name, latent=None, denoise=1.0, tiled=False, cfg_scale=7.0):
|
||||
def vae_parse(x, tiled=False):
|
||||
if final_vae is None or final_refiner_vae is None:
|
||||
return x
|
||||
|
||||
print('VAE parsing ...')
|
||||
x = core.decode_vae(vae=final_vae, latent_image=x, tiled=tiled)
|
||||
x = core.encode_vae(vae=final_refiner_vae, pixels=x, tiled=tiled)
|
||||
print('VAE parsed ...')
|
||||
|
||||
return x
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
@torch.inference_mode()
|
||||
def process_diffusion(positive_cond, negative_cond, steps, switch, width, height, image_seed, callback, sampler_name, scheduler_name, latent=None, denoise=1.0, tiled=False, cfg_scale=7.0, refiner_swap_method='joint'):
|
||||
assert refiner_swap_method in ['joint', 'separate', 'vae']
|
||||
print(f'[Sampler] refiner_swap_method = {refiner_swap_method}')
|
||||
|
||||
if latent is None:
|
||||
empty_latent = core.generate_empty_latent(width=width, height=height, batch_size=1)
|
||||
else:
|
||||
empty_latent = latent
|
||||
|
||||
sampled_latent = core.ksampler(
|
||||
model=final_unet,
|
||||
refiner=final_refiner,
|
||||
positive=positive_cond,
|
||||
negative=negative_cond,
|
||||
latent=empty_latent,
|
||||
steps=steps, start_step=0, last_step=steps, disable_noise=False, force_full_denoise=True,
|
||||
seed=image_seed,
|
||||
denoise=denoise,
|
||||
callback_function=callback,
|
||||
cfg=cfg_scale,
|
||||
sampler_name=sampler_name,
|
||||
scheduler=scheduler_name,
|
||||
refiner_switch=switch
|
||||
)
|
||||
if refiner_swap_method == 'joint':
|
||||
sampled_latent = core.ksampler(
|
||||
model=final_unet,
|
||||
refiner=final_refiner_unet,
|
||||
positive=positive_cond,
|
||||
negative=negative_cond,
|
||||
latent=empty_latent,
|
||||
steps=steps, start_step=0, last_step=steps, disable_noise=False, force_full_denoise=True,
|
||||
seed=image_seed,
|
||||
denoise=denoise,
|
||||
callback_function=callback,
|
||||
cfg=cfg_scale,
|
||||
sampler_name=sampler_name,
|
||||
scheduler=scheduler_name,
|
||||
refiner_switch=switch,
|
||||
previewer_start=0,
|
||||
previewer_end=steps,
|
||||
)
|
||||
decoded_latent = core.decode_vae(vae=final_vae, latent_image=sampled_latent, tiled=tiled)
|
||||
images = core.pytorch_to_numpy(decoded_latent)
|
||||
return images
|
||||
|
||||
decoded_latent = core.decode_vae(vae=final_vae, latent_image=sampled_latent, tiled=tiled)
|
||||
images = core.pytorch_to_numpy(decoded_latent)
|
||||
if refiner_swap_method == 'separate':
|
||||
sampled_latent = core.ksampler(
|
||||
model=final_unet,
|
||||
positive=positive_cond,
|
||||
negative=negative_cond,
|
||||
latent=empty_latent,
|
||||
steps=steps, start_step=0, last_step=switch, disable_noise=False, force_full_denoise=False,
|
||||
seed=image_seed,
|
||||
denoise=denoise,
|
||||
callback_function=callback,
|
||||
cfg=cfg_scale,
|
||||
sampler_name=sampler_name,
|
||||
scheduler=scheduler_name,
|
||||
previewer_start=0,
|
||||
previewer_end=switch,
|
||||
)
|
||||
print('Refiner swapped by changing ksampler. Noise preserved.')
|
||||
|
||||
comfy.model_management.soft_empty_cache()
|
||||
return images
|
||||
target_model = final_refiner_unet
|
||||
if target_model is None:
|
||||
target_model = final_unet
|
||||
print('Use base model to refine itself - this may because of developer mode.')
|
||||
|
||||
sampled_latent = core.ksampler(
|
||||
model=target_model,
|
||||
positive=clip_separate(positive_cond, target_model=target_model.model),
|
||||
negative=clip_separate(negative_cond, target_model=target_model.model),
|
||||
latent=sampled_latent,
|
||||
steps=steps, start_step=switch, last_step=steps, disable_noise=True, force_full_denoise=True,
|
||||
seed=image_seed,
|
||||
denoise=denoise,
|
||||
callback_function=callback,
|
||||
cfg=cfg_scale,
|
||||
sampler_name=sampler_name,
|
||||
scheduler=scheduler_name,
|
||||
previewer_start=switch,
|
||||
previewer_end=steps,
|
||||
)
|
||||
|
||||
target_model = final_refiner_vae
|
||||
if target_model is None:
|
||||
target_model = final_vae
|
||||
decoded_latent = core.decode_vae(vae=target_model, latent_image=sampled_latent, tiled=tiled)
|
||||
images = core.pytorch_to_numpy(decoded_latent)
|
||||
return images
|
||||
|
||||
if refiner_swap_method == 'vae':
|
||||
sampled_latent = core.ksampler(
|
||||
model=final_unet,
|
||||
positive=positive_cond,
|
||||
negative=negative_cond,
|
||||
latent=empty_latent,
|
||||
steps=steps, start_step=0, last_step=switch, disable_noise=False, force_full_denoise=True,
|
||||
seed=image_seed,
|
||||
denoise=denoise,
|
||||
callback_function=callback,
|
||||
cfg=cfg_scale,
|
||||
sampler_name=sampler_name,
|
||||
scheduler=scheduler_name,
|
||||
previewer_start=0,
|
||||
previewer_end=switch,
|
||||
)
|
||||
print('Refiner swapped by changing ksampler. Noise is not preserved.')
|
||||
|
||||
target_model = final_refiner_unet
|
||||
if target_model is None:
|
||||
target_model = final_unet
|
||||
print('Use base model to refine itself - this may because of developer mode.')
|
||||
|
||||
sampled_latent = vae_parse(sampled_latent)
|
||||
|
||||
sampled_latent = core.ksampler(
|
||||
model=target_model,
|
||||
positive=clip_separate(positive_cond, target_model=target_model.model),
|
||||
negative=clip_separate(negative_cond, target_model=target_model.model),
|
||||
latent=sampled_latent,
|
||||
steps=steps, start_step=switch, last_step=steps, disable_noise=False, force_full_denoise=True,
|
||||
seed=image_seed,
|
||||
denoise=denoise,
|
||||
callback_function=callback,
|
||||
cfg=cfg_scale,
|
||||
sampler_name=sampler_name,
|
||||
scheduler=scheduler_name,
|
||||
previewer_start=switch,
|
||||
previewer_end=steps,
|
||||
)
|
||||
|
||||
target_model = final_refiner_vae
|
||||
if target_model is None:
|
||||
target_model = final_vae
|
||||
decoded_latent = core.decode_vae(vae=target_model, latent_image=sampled_latent, tiled=tiled)
|
||||
images = core.pytorch_to_numpy(decoded_latent)
|
||||
return images
|
||||
|
||||
Reference in New Issue
Block a user