mirror of https://github.com/vladmandic/automatic
77 lines
3.6 KiB
Python
77 lines
3.6 KiB
Python
import os
|
|
import huggingface_hub as hf
|
|
from modules import shared, processing, sd_models, devices
|
|
|
|
|
|
def photo_maker(p: processing.StableDiffusionProcessing, input_images, trigger, strength, start): # pylint: disable=arguments-differ
|
|
from modules.face.photomaker_model import PhotoMakerStableDiffusionXLPipeline
|
|
|
|
# prepare pipeline
|
|
if len(input_images) == 0:
|
|
shared.log.warning('PhotoMaker: no input images')
|
|
return None
|
|
|
|
c = shared.sd_model.__class__.__name__ if shared.sd_loaded else ''
|
|
if c != 'StableDiffusionXLPipeline':
|
|
shared.log.warning(f'PhotoMaker invalid base model: current={c} required=StableDiffusionXLPipeline')
|
|
return None
|
|
|
|
# validate prompt
|
|
if p.all_prompts is None or len(p.all_prompts) == 0:
|
|
processing.process_init(p)
|
|
p.init(p.all_prompts, p.all_seeds, p.all_subseeds)
|
|
trigger_ids = shared.sd_model.tokenizer.encode(trigger) + shared.sd_model.tokenizer_2.encode(trigger)
|
|
prompt_ids1 = shared.sd_model.tokenizer.encode(p.all_prompts[0])
|
|
prompt_ids2 = shared.sd_model.tokenizer_2.encode(p.all_prompts[0])
|
|
for t in trigger_ids:
|
|
if prompt_ids1.count(t) != 1:
|
|
shared.log.error(f'PhotoMaker: trigger word not matched in prompt: {trigger} ids={trigger_ids} prompt={p.all_prompts[0]} ids={prompt_ids1}')
|
|
return None
|
|
if prompt_ids2.count(t) != 1:
|
|
shared.log.error(f'PhotoMaker: trigger word not matched in prompt: {trigger} ids={trigger_ids} prompt={p.all_prompts[0]} ids={prompt_ids1}')
|
|
return None
|
|
|
|
# create new pipeline
|
|
orig_pipeline = shared.sd_model # backup current pipeline definition
|
|
shared.sd_model = PhotoMakerStableDiffusionXLPipeline(
|
|
vae = shared.sd_model.vae,
|
|
text_encoder=shared.sd_model.text_encoder,
|
|
text_encoder_2=shared.sd_model.text_encoder_2,
|
|
tokenizer=shared.sd_model.tokenizer,
|
|
tokenizer_2=shared.sd_model.tokenizer_2,
|
|
unet=shared.sd_model.unet,
|
|
scheduler=shared.sd_model.scheduler,
|
|
force_zeros_for_empty_prompt=shared.opts.diffusers_force_zeros,
|
|
)
|
|
sd_models.copy_diffuser_options(shared.sd_model, orig_pipeline) # copy options from original pipeline
|
|
sd_models.set_diffuser_options(shared.sd_model) # set all model options such as fp16, offload, etc.
|
|
sd_models.move_model(shared.sd_model, devices.device) # move pipeline to device
|
|
shared.sd_model.to(dtype=devices.dtype)
|
|
|
|
orig_prompt_attention = shared.opts.prompt_attention
|
|
shared.opts.data['prompt_attention'] = 'Fixed attention' # otherwise need to deal with class_tokens_mask
|
|
p.task_args['input_id_images'] = input_images
|
|
p.task_args['start_merge_step'] = int(start * p.steps)
|
|
p.task_args['prompt'] = p.all_prompts[0] if p.all_prompts is not None else p.prompt
|
|
|
|
photomaker_path = hf.hf_hub_download(repo_id="TencentARC/PhotoMaker", filename="photomaker-v1.bin", repo_type="model", cache_dir=shared.opts.diffusers_dir)
|
|
shared.log.debug(f'PhotoMaker: model={photomaker_path} images={len(input_images)} trigger={trigger} args={p.task_args}')
|
|
|
|
# load photomaker adapter
|
|
shared.sd_model.load_photomaker_adapter(
|
|
os.path.dirname(photomaker_path),
|
|
subfolder="",
|
|
weight_name=os.path.basename(photomaker_path),
|
|
trigger_word=trigger
|
|
)
|
|
shared.sd_model.set_adapters(["photomaker"], adapter_weights=[strength])
|
|
|
|
# run processing
|
|
processed: processing.Processed = processing.process_images(p)
|
|
p.extra_generation_params['PhotoMaker'] = f'{strength}'
|
|
|
|
# restore original pipeline
|
|
shared.opts.data['prompt_attention'] = orig_prompt_attention
|
|
shared.sd_model = orig_pipeline
|
|
return processed
|