mirror of
https://github.com/vladmandic/sdnext.git
synced 2026-01-27 15:02:48 +03:00
122 lines
5.2 KiB
Python
122 lines
5.2 KiB
Python
# https://huggingface.co/ByteDance/InfiniteYou
|
|
# https://github.com/bytedance/InfiniteYou
|
|
# flux base model + 11.8gb controlnet module + 338mb image module + 428 insightface module
|
|
|
|
import gradio as gr
|
|
from PIL import Image
|
|
from modules import scripts_manager, processing, shared, sd_models, devices
|
|
|
|
|
|
prefix = 'InfiniteYou'
|
|
model_versions = ['aes_stage2', 'sim_stage1']
|
|
orig_pipeline, orig_prompt_attention = None, None
|
|
|
|
|
|
def verify_insightface():
|
|
from installer import installed, install_insightface
|
|
if not installed('insightface', reload=False, quiet=True):
|
|
install_insightface()
|
|
|
|
|
|
def load_infiniteyou(model: str):
|
|
from scripts.infiniteyou import InfUFluxPipeline
|
|
shared.sd_model = InfUFluxPipeline(
|
|
pipe=shared.sd_model,
|
|
model_version=model,
|
|
)
|
|
sd_models.copy_diffuser_options(shared.sd_model, orig_pipeline)
|
|
sd_models.clear_caches(full=True)
|
|
sd_models.set_diffuser_options(shared.sd_model)
|
|
|
|
|
|
class Script(scripts_manager.Script):
|
|
def title(self):
|
|
return f'{prefix}: Flexible Photo Recrafting'
|
|
|
|
def show(self, is_img2img):
|
|
return not is_img2img
|
|
|
|
# return signature is array of gradio components
|
|
def ui(self, _is_img2img):
|
|
with gr.Row():
|
|
gr.HTML(f'<a href="https://github.com/bytedance/InfiniteYou">  {prefix}: Flexible Photo Recrafting</a><br>')
|
|
with gr.Row():
|
|
model = gr.Dropdown(label='IY model', choices=model_versions, value=model_versions[0])
|
|
restore = gr.Checkbox(label='Restore pipeline on end', value=False)
|
|
with gr.Row():
|
|
scale = gr.Slider(label='IY scale', value=1.0, minimum=0.0, maximum=2.0, step=0.05)
|
|
with gr.Row():
|
|
start = gr.Slider(label='IY start', value=0.0, minimum=0.0, maximum=1.0, step=0.05)
|
|
end = gr.Slider(label='IY end', value=1.0, minimum=0.0, maximum=1.0, step=0.05)
|
|
with gr.Row():
|
|
id_guidance = gr.Slider(label='Identity guidance', value=3.5, minimum=0.0, maximum=14.0, step=0.05)
|
|
with gr.Row():
|
|
id_image = gr.Image(label='Identity image', type='pil')
|
|
with gr.Row():
|
|
control_guidance = gr.Slider(label='Control guidance', value=1.0, minimum=0.0, maximum=14.0, step=0.05)
|
|
with gr.Row():
|
|
control_image = gr.Image(label='Control image', type='pil')
|
|
return [model, id_image, control_image, scale, start, end, id_guidance, control_guidance, restore]
|
|
|
|
def run(self, p: processing.StableDiffusionProcessing,
|
|
model: str = None,
|
|
id_image: Image.Image = None,
|
|
control_image: Image.Image = None,
|
|
scale: float = 1.0,
|
|
start: float = 0.0,
|
|
end: float = 1.0,
|
|
id_guidance: float = 3.5,
|
|
control_guidance: float = 1.0,
|
|
restore: bool = False,
|
|
): # pylint: disable=arguments-differ, unused-argument
|
|
|
|
if model is None or model not in model_versions:
|
|
return None
|
|
if id_image is None:
|
|
shared.log.error(f'{prefix}: no init_images')
|
|
return None
|
|
if shared.sd_model_type != 'f1':
|
|
shared.log.error(f'{prefix}: invalid model type: {shared.sd_model_type}')
|
|
return None
|
|
if scale <= 0:
|
|
return None
|
|
|
|
global orig_pipeline, orig_prompt_attention # pylint: disable=global-statement
|
|
orig_pipeline = shared.sd_model
|
|
if shared.sd_model.__class__.__name__ != 'InfUFluxPipeline':
|
|
verify_insightface()
|
|
load_infiniteyou(model)
|
|
devices.torch_gc()
|
|
shared.log.info(f'{prefix}: cls={shared.sd_model.__class__.__name__} loaded')
|
|
|
|
processing.fix_seed(p)
|
|
p.task_args['id_image'] = id_image
|
|
p.task_args['control_image'] = control_image
|
|
p.task_args['infusenet_conditioning_scale'] = p.task_args.get('infusenet_conditioning_scale', scale)
|
|
p.task_args['infusenet_guidance_start'] = p.task_args.get('infusenet_guidance_start', start)
|
|
p.task_args['infusenet_guidance_end'] = p.task_args.get('infusenet_guidance_end', end)
|
|
p.task_args['seed'] = p.seed
|
|
p.task_args['negative_prompt'] = None
|
|
p.task_args['guidance_scale'] = id_guidance
|
|
p.task_args['controlnet_guidance_scale'] = control_guidance
|
|
p.extra_generation_params['IY model'] = model
|
|
p.extra_generation_params['IY guidance'] = f'{scale:.1f}/{start:.1f}/{end:.1f}'
|
|
orig_prompt_attention = shared.opts.prompt_attention
|
|
shared.opts.data['prompt_attention'] = 'fixed'
|
|
shared.log.debug(f'{prefix}: args={p.task_args}')
|
|
|
|
processed = processing.process_images(p)
|
|
return processed
|
|
|
|
def after(self, p: processing.StableDiffusionProcessing, processed: processing.Processed, *args, **kwargs): # pylint: disable=unused-argument
|
|
# restore pipeline
|
|
global orig_pipeline, orig_prompt_attention # pylint: disable=global-statement
|
|
restore = args[-1]
|
|
if orig_prompt_attention is not None:
|
|
shared.opts.data['prompt_attention'] = orig_prompt_attention
|
|
orig_prompt_attention = None
|
|
if restore and orig_pipeline is not None:
|
|
shared.log.info(f'{prefix}: restoring pipeline')
|
|
shared.sd_model = orig_pipeline
|
|
orig_pipeline = None
|