mirror of
https://github.com/vladmandic/automatic
synced 2026-09-19 17:24:32 +02:00
diffusers code refactoring and exception handling
This commit is contained in:
Vladimir Mandic
+2
-136
@@ -3,7 +3,6 @@ import math
|
||||
import os
|
||||
import hashlib
|
||||
import random
|
||||
import inspect
|
||||
from contextlib import nullcontext
|
||||
from typing import Any, Dict, List
|
||||
import torch
|
||||
@@ -26,7 +25,6 @@ import modules.images as images
|
||||
import modules.styles
|
||||
import modules.sd_models as sd_models
|
||||
import modules.sd_vae as sd_vae
|
||||
from modules.lora_diffusers import lora_state, unload_diffusers_lora
|
||||
|
||||
|
||||
opt_C = 4
|
||||
@@ -436,7 +434,6 @@ def decode_first_stage(model, x):
|
||||
def get_fixed_seed(seed):
|
||||
if seed is None or seed == '' or seed == -1:
|
||||
return int(random.randrange(4294967294))
|
||||
|
||||
return seed
|
||||
|
||||
|
||||
@@ -623,53 +620,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
|
||||
cache[0] = (required_prompts, steps)
|
||||
return cache[1]
|
||||
|
||||
# TODO Diffusers limited callbacks
|
||||
def diffusers_callback(step: int, _timestep: int, latents: torch.FloatTensor):
|
||||
shared.state.sampling_step = step
|
||||
shared.state.sampling_steps = p.steps
|
||||
shared.state.current_latent = latents
|
||||
|
||||
def set_pipeline_args(model, prompt, negative_prompt, **kwargs):
|
||||
args = {}
|
||||
pipeline = model.main if model.__class__.__name__ == 'PriorPipeline' else model
|
||||
signature = inspect.signature(type(pipeline).__call__)
|
||||
possible = signature.parameters.keys()
|
||||
generator_device = 'cpu' if shared.opts.diffusers_generator_device == "cpu" else shared.device
|
||||
generator = [torch.Generator(generator_device).manual_seed(s) for s in seeds]
|
||||
if 'prompt' in possible:
|
||||
args['prompt'] = prompt
|
||||
if 'negative_prompt' in possible:
|
||||
args['negative_prompt'] = negative_prompt
|
||||
if 'num_inference_steps' in possible:
|
||||
args['num_inference_steps'] = p.steps
|
||||
if 'guidance_scale' in possible:
|
||||
args['guidance_scale'] = p.cfg_scale
|
||||
if 'generator' in possible:
|
||||
args['generator'] = generator
|
||||
if 'output_type' in possible:
|
||||
args['output_type'] = 'np'
|
||||
if 'callback_steps' in possible:
|
||||
args['callback_steps'] = 1
|
||||
if 'callback' in possible:
|
||||
args['callback'] = diffusers_callback
|
||||
if 'cross_attention_kwargs' in possible:
|
||||
args['cross_attention_kwargs'] = cross_attention_kwargs
|
||||
for arg in kwargs:
|
||||
if arg in possible:
|
||||
args[arg] = kwargs[arg]
|
||||
# shared.log.debug(f'Diffuser pipeline: {pipeline.__class__.__name__} possible={possible}')
|
||||
clean = args.copy()
|
||||
clean.pop('callback', None)
|
||||
clean.pop('callback_steps', None)
|
||||
clean.pop('image', None)
|
||||
clean.pop('mask_image', None)
|
||||
clean.pop('prompt', None)
|
||||
clean.pop('negative_prompt', None)
|
||||
clean['generator'] = generator_device
|
||||
shared.log.debug(f'Diffuser pipeline: {pipeline.__class__.__name__} set={clean}')
|
||||
return args
|
||||
|
||||
|
||||
ema_scope_context = p.sd_model.ema_scope if shared.backend == shared.Backend.ORIGINAL else nullcontext
|
||||
with torch.no_grad(), ema_scope_context():
|
||||
with devices.autocast():
|
||||
@@ -740,92 +690,8 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
|
||||
del samples_ddim
|
||||
|
||||
elif shared.backend == shared.Backend.DIFFUSERS:
|
||||
if (not hasattr(shared.sd_model.scheduler, 'name')) or (shared.sd_model.scheduler.name != p.sampler_name) and (p.sampler_name != 'Default'):
|
||||
sampler = sd_samplers.all_samplers_map.get(p.sampler_name, None)
|
||||
if sampler is None:
|
||||
sampler = sd_samplers.all_samplers_map.get("UniPC")
|
||||
sd_samplers.create_sampler(sampler.name, shared.sd_model) # TODO(Patrick): For wrapped pipelines this is currently a no-op
|
||||
|
||||
cross_attention_kwargs={}
|
||||
if lora_state['active']:
|
||||
cross_attention_kwargs['scale'] = lora_state['multiplier']
|
||||
task_specific_kwargs={}
|
||||
if sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.TEXT_2_IMAGE:
|
||||
task_specific_kwargs = {"height": p.height, "width": p.width}
|
||||
elif sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.IMAGE_2_IMAGE:
|
||||
task_specific_kwargs = {"image": p.init_images[0], "strength": p.denoising_strength}
|
||||
elif sd_models.get_diffusers_task(shared.sd_model) == sd_models.DiffusersTaskType.INPAINTING:
|
||||
# TODO(PVP): change out to latents once possible with `diffusers`
|
||||
task_specific_kwargs = {"image": p.init_images[0], "mask_image": p.image_mask, "strength": p.denoising_strength}
|
||||
|
||||
shared.sd_model.to(devices.device)
|
||||
pipe_args = set_pipeline_args(
|
||||
model=shared.sd_model,
|
||||
prompt=prompts,
|
||||
negative_prompt=negative_prompts,
|
||||
eta=shared.opts.eta_ddim,
|
||||
guidance_rescale=p.diffusers_guidance_rescale,
|
||||
# aesthetic_score=shared.opts.diffusers_aesthetics_score,
|
||||
output_type='np' if (shared.sd_refiner is None or p.enable_hr is False or not shared.opts.diffusers_refiner_latents) else 'latent',
|
||||
**task_specific_kwargs
|
||||
)
|
||||
output = shared.sd_model(**pipe_args) # pylint: disable=not-callable
|
||||
|
||||
# TODO Diffusers processing is not using p.sample so second pass is ignored and we use this instead
|
||||
if shared.sd_refiner is not None and p.enable_hr:
|
||||
if shared.opts.diffusers_move_base:
|
||||
shared.log.debug('Moving base model to CPU')
|
||||
shared.sd_model.to('cpu')
|
||||
|
||||
if (not hasattr(shared.sd_refiner.scheduler, 'name')) or (shared.sd_refiner.scheduler.name != p.latent_sampler) and (p.sampler_name != 'Default'):
|
||||
sampler = sd_samplers.all_samplers_map.get(p.latent_sampler, None)
|
||||
if sampler is None:
|
||||
sampler = sd_samplers.all_samplers_map.get("UniPC")
|
||||
sd_samplers.create_sampler(sampler.name, shared.sd_refiner) # TODO(Patrick): For wrapped pipelines this is currently a no-op
|
||||
|
||||
shared.sd_refiner.to(devices.device)
|
||||
devices.torch_gc()
|
||||
x_samples_ddim = []
|
||||
|
||||
for i in range(len(output.images)):
|
||||
|
||||
"""
|
||||
# TODO save before refiner
|
||||
if shared.opts.save and not p.do_not_save_samples and shared.opts.save_images_before_refiner and hasattr(shared.sd_model, 'vae'):
|
||||
info=infotext(n, i)
|
||||
image = decode_first_stage(shared.sd_model, output.images[i].to(dtype=devices.dtype_vae))
|
||||
images.save_image(image, path=p.outpath_samples, basename="", seed=seeds[i], prompt=prompts[i], extension=shared.opts.samples_format, info=info, p=p, suffix="-before-refiner")
|
||||
"""
|
||||
|
||||
pipe_args = set_pipeline_args(
|
||||
model=shared.sd_refiner,
|
||||
prompt=[p.refiner_prompt] if len(p.refiner_prompt) > 0 else prompts,
|
||||
negative_prompt=[p.refiner_negative] if len(p.refiner_negative) > 0 else negative_prompts,
|
||||
num_inference_steps=p.hr_second_pass_steps,
|
||||
eta=shared.opts.eta_ddim,
|
||||
strength=p.denoising_strength,
|
||||
guidance_scale=p.image_cfg_scale if p.image_cfg_scale is not None else p.cfg_scale,
|
||||
guidance_rescale=p.diffusers_guidance_rescale,
|
||||
# aesthetic_score=shared.opts.diffusers_aesthetics_score,
|
||||
denoising_start=p.refiner_denoise_start,
|
||||
denoising_end=p.refiner_denoise_end,
|
||||
image=output.images[i],
|
||||
output_type='np',
|
||||
)
|
||||
output = shared.sd_refiner(**pipe_args) # pylint: disable=not-callable
|
||||
x_samples_ddim.append(output.images[0])
|
||||
|
||||
if shared.opts.diffusers_move_refiner:
|
||||
shared.log.debug('Moving refiner model to CPU')
|
||||
shared.sd_refiner.to('cpu')
|
||||
else:
|
||||
x_samples_ddim = output.images
|
||||
|
||||
if p.is_hr_pass:
|
||||
shared.log.warning('Diffusers not implemented: hires fix')
|
||||
|
||||
if lora_state['active']:
|
||||
unload_diffusers_lora()
|
||||
from modules.processing_diffusers import process_diffusers
|
||||
x_samples_ddim = process_diffusers(p, seeds, prompts, negative_prompts)
|
||||
|
||||
else:
|
||||
raise ValueError(f"Unknown backend {shared.backend}")
|
||||
|
||||
Reference in New Issue
Block a user