# -*- coding: utf-8 -*- # Copyright (c) Alibaba, Inc. and its affiliates. import copy import os.path import random from collections import OrderedDict import gradio as gr import torch import torchvision.transforms.functional as TF from scepter.modules.model.network.diffusion.diffusion import GaussianDiffusion from scepter.modules.model.utils.data_utils import crop_back from scepter.modules.utils.distribute import we from scepter.modules.utils.file_system import FS from .diffusion_inference import DiffusionInference, get_model class LargenInference(DiffusionInference): ''' define vae, unet, text-encoder, tuner, refiner components support to load the components dynamicly. create and load model when run this model at the first time. ''' def __init__(self, logger=None): self.logger = logger self.is_redefine_paras = True self.loaded_model = {} self.loaded_model_name = [ 'diffusion_model', 'first_stage_model', 'cond_stage_model' ] self.diffusion_insclass = GaussianDiffusion def redefine_paras(self, cfg): if cfg.get('PRETRAINED_MODEL', None): assert FS.isfile(cfg.PRETRAINED_MODEL) with FS.get_from(cfg.PRETRAINED_MODEL, wait_finish=True) as local_path: if local_path.endswith('safetensors'): from safetensors.torch import load_file as load_safetensors sd = load_safetensors(local_path) else: if 'weights_only' in torch.load.__code__.co_varnames: sd = torch.load(local_path, map_location='cpu', weights_only=True) else: sd = torch.load(local_path, map_location='cpu') if 'model' in sd: sd = sd['model'] first_stage_model_path = os.path.join( os.path.dirname(local_path), 'first_stage_model.pth') cond_stage_model_path = os.path.join( os.path.dirname(local_path), 'cond_stage_model.pth') diffusion_model_path = os.path.join( os.path.dirname(local_path), 'diffusion_model.pth') if (not os.path.exists(first_stage_model_path) or not os.path.exists(cond_stage_model_path) or not os.path.exists(diffusion_model_path)): self.logger.info( 'Now read the whole model and rearrange the modules, it may take several mins.' ) first_stage_model = OrderedDict() cond_stage_model = OrderedDict() diffusion_model = OrderedDict() for k, v in sd.items(): if k.startswith('first_stage_model.'): first_stage_model[k.replace( 'first_stage_model.', '')] = v elif k.startswith('conditioner.'): cond_stage_model[k.replace('conditioner.', '')] = v elif k.startswith('cond_stage_model.'): if k.startswith('cond_stage_model.model.'): cond_stage_model[k.replace( 'cond_stage_model.model.', '')] = v else: cond_stage_model[k.replace( 'cond_stage_model.', '')] = v elif k.startswith('model.diffusion_model.'): diffusion_model[k.replace('model.diffusion_model.', '')] = v elif k.startswith('model.'): diffusion_model[k.replace('model.', '')] = v else: continue if cfg.have('FIRST_STAGE_MODEL'): with open(first_stage_model_path + 'cache', 'wb') as f: torch.save(first_stage_model, f) os.rename(first_stage_model_path + 'cache', first_stage_model_path) self.logger.info( 'First stage model has been processed.') if cfg.have('COND_STAGE_MODEL'): with open(cond_stage_model_path + 'cache', 'wb') as f: torch.save(cond_stage_model, f) os.rename(cond_stage_model_path + 'cache', cond_stage_model_path) self.logger.info( 'Cond stage model has been processed.') if cfg.have('DIFFUSION_MODEL'): with open(diffusion_model_path + 'cache', 'wb') as f: torch.save(diffusion_model, f) os.rename(diffusion_model_path + 'cache', diffusion_model_path) self.logger.info('Diffusion model has been processed.') if not cfg.FIRST_STAGE_MODEL.get('PRETRAINED_MODEL', None): cfg.FIRST_STAGE_MODEL.PRETRAINED_MODEL = first_stage_model_path else: cfg.FIRST_STAGE_MODEL.RELOAD_MODEL = first_stage_model_path if not cfg.COND_STAGE_MODEL.get('PRETRAINED_MODEL', None): cfg.COND_STAGE_MODEL.PRETRAINED_MODEL = cond_stage_model_path else: cfg.COND_STAGE_MODEL.RELOAD_MODEL = cond_stage_model_path if not cfg.DIFFUSION_MODEL.get('PRETRAINED_MODEL', None): cfg.DIFFUSION_MODEL.PRETRAINED_MODEL = diffusion_model_path else: cfg.DIFFUSION_MODEL.RELOAD_MODEL = diffusion_model_path return cfg @torch.no_grad() def __call__(self, input, num_samples=1, intermediate_callback=None, refine_strength=0, img_to_img_strength=0, cat_uc=True, tuner_model=None, control_model=None, largen_state=False, **kwargs): if not largen_state: raise gr.Error('LARGEN model must be used with LAR-Gen settings') value_input = copy.deepcopy(self.input) value_input.update(input) print(value_input) height, width = value_input['target_size_as_tuple'] value_output = copy.deepcopy(self.output) batch, batch_uc = self.get_batch(value_input, num_samples=1) # first stage encode task = kwargs.get('largen_task', 'Text_Guided_Inpainting') image_scale = kwargs.get('largen_image_scale', 1.0) tar_image = kwargs.get('largen_tar_image', None) tar_mask = kwargs.get('largen_tar_mask', None) masked_image = kwargs.get('largen_masked_image', None) ref_image = kwargs.get('largen_ref_image', None) ref_mask = kwargs.get('largen_ref_mask', None) ref_clip = kwargs.get('largen_ref_clip', None) base_image = kwargs.get('largen_base_image', None) extra_sizes = kwargs.get('largen_extra_sizes', None) bbox_yyxx = kwargs.get('largen_bbox_yyxx', None) device = we.device_id tar_image = tar_image.to(device) tar_mask = tar_mask.to(device) masked_image = masked_image.to(device) if 'Subject' in task: ref_image = ref_image.to(device) ref_mask = ref_mask.to(device) ref_clip = ref_clip.to(device) self.dynamic_load(self.first_stage_model, 'first_stage_model') tar_x0 = self.encode_first_stage(tar_image) masked_x0 = self.encode_first_stage(masked_image) b, _, h, w = tar_x0.shape tar_mask_latent = TF.resize(tar_mask, (h, w), antialias=True) tar_mask_latent = (tar_mask_latent > 0.5).float() batch.update({ 'tar_x0': tar_x0, 'tar_mask_latent': tar_mask_latent, 'masked_x0': masked_x0, 'task': task }) batch_uc.update({ 'tar_x0': tar_x0, 'tar_mask_latent': tar_mask_latent, 'masked_x0': masked_x0, 'task': task }) if 'Subject' in task and ref_image is not None: ref_x0 = self.encode_first_stage(ref_image) batch.update({ 'ref_ip': ref_clip, 'ref_detail': ref_clip, 'ref_x0': ref_x0, 'ref_mask': ref_mask, 'image_scale': image_scale, }) batch_uc.update({ 'ref_ip': torch.zeros_like(ref_clip), 'ref_detail': ref_clip, 'ref_x0': ref_x0, 'ref_mask': ref_mask, 'image_scale': image_scale, }) self.dynamic_unload(self.first_stage_model, 'first_stage_model', skip_loaded=False) # cond stage self.dynamic_load(self.cond_stage_model, 'cond_stage_model') function_name, dtype = self.get_function_info(self.cond_stage_model) with torch.autocast('cuda', enabled=dtype == 'float16', dtype=getattr(torch, dtype)): context = getattr(get_model(self.cond_stage_model), function_name)(batch) null_context = getattr(get_model(self.cond_stage_model), function_name)(batch_uc) self.dynamic_unload(self.cond_stage_model, 'cond_stage_model', skip_loaded=False) # get noise seed = kwargs.pop('seed', -1) g = torch.Generator(device=we.device_id) seed = seed if seed >= 0 else random.randint(0, 2**32 - 1) g.manual_seed(seed) if 'seed' in value_output: value_output['seed'] = seed for sample_id in range(num_samples): if self.diffusion_model is not None: noise = torch.empty( 1, 4, height // self.first_stage_model['paras']['size_factor'], width // self.first_stage_model['paras']['size_factor'], device=we.device_id).normal_(generator=g) self.dynamic_load(self.diffusion_model, 'diffusion_model') # UNet use input n_prompt function_name, dtype = self.get_function_info( self.diffusion_model) with torch.autocast('cuda', enabled=dtype == 'float16', dtype=getattr(torch, dtype)): latent = self.diffusion.sample( noise=noise, x=None, denoising_strength=1.0, refine_strength=refine_strength, solver=value_input.get('sample', 'ddim'), model=get_model(self.diffusion_model), model_kwargs=[{ 'cond': context }, { 'cond': null_context }], steps=value_input.get('sample_steps', 50), guide_scale=value_input.get('guide_scale', 7.5), guide_rescale=value_input.get('guide_rescale', 0.5), discretization=value_input.get('discretization', 'trailing'), show_progress=True, seed=seed, condition_fn=None, clamp=None, percentile=None, t_max=None, t_min=None, discard_penultimate_step=None, intermediate_callback=intermediate_callback, cat_uc=cat_uc, **kwargs) self.dynamic_unload(self.diffusion_model, 'diffusion_model', skip_loaded=True) if 'latent' in value_output: if value_output['latent'] is None or ( isinstance(value_output['latent'], list) and len(value_output['latent']) < 1): value_output['latent'] = [] value_output['latent'].append(latent) self.dynamic_load(self.first_stage_model, 'first_stage_model') x_samples = self.decode_first_stage(latent).float() self.dynamic_unload(self.first_stage_model, 'first_stage_model', skip_loaded=False) images = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0) if base_image is not None: stitch_images = [] for img in images: stitch_img = crop_back(img, copy.deepcopy(base_image), extra_sizes, bbox_yyxx) stitch_images.append(stitch_img) images = torch.stack(stitch_images, dim=0) if 'images' in value_output: if value_output['images'] is None or ( isinstance(value_output['images'], list) and len(value_output['images']) < 1): value_output['images'] = [] value_output['images'].append(images) for k, v in value_output.items(): if isinstance(v, list): value_output[k] = torch.cat(v, dim=0) if isinstance(v, torch.Tensor): value_output[k] = v.cpu() return value_output