diff --git a/hellomeme/pipelines/pipline_hm_image.py b/hellomeme/pipelines/pipline_hm_image.py index 56275aa..34f4f6b 100644 --- a/hellomeme/pipelines/pipline_hm_image.py +++ b/hellomeme/pipelines/pipline_hm_image.py @@ -27,6 +27,11 @@ from ..models import HMReferenceAdapter class HMImagePipeline(StableDiffusionImg2ImgPipeline): def caryomitosis(self, **kwargs): + if hasattr(self, "unet_ref"): + del self.unet_ref + self.unet_ref = HMDenoising3D.from_unet2d(self.unet) + self.unet_ref.cpu() + if not isinstance(self.unet, HMDenoising3D): unet = HMDenoising3D.from_unet2d(unet=self.unet) # todo: 不够优雅 @@ -34,6 +39,8 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline): self.unet = unet self.unet.cpu() + self.vae_decode = copy.deepcopy(self.vae) + def insert_hm_modules(self, version, dtype): if isinstance(self.unet, HMDenoising3D): if version == 'v1': @@ -43,6 +50,9 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline): self.unet.insert_reference_adapter(hm_adapter) self.unet.to(device='cpu', dtype=dtype).eval() + if hasattr(self, "unet_ref"): + self.unet_ref.to(device='cpu', dtype=dtype).eval() + if hasattr(self, "mp_control"): del self.mp_control if version == 'v1': @@ -60,6 +70,7 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline): self.mp_control2.to(device='cpu', dtype=dtype).eval() self.vae.to(device='cpu', dtype=dtype).eval() + self.vae_decode.to(device='cpu', dtype=dtype).eval() self.text_encoder.to(device='cpu', dtype=dtype).eval() @torch.no_grad() @@ -232,13 +243,14 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline): ) latent_model_input = torch.cat([torch.zeros_like(ref_latents), ref_latents]) if self.do_classifier_free_guidance else ref_latents - self.unet.to(device=device) - cached_res = self.unet( + self.unet_ref.to(device=device) + cached_res = self.unet_ref( latent_model_input.unsqueeze(2), 0, encoder_hidden_states=prompt_embeds, return_dict=False, )[1] + self.unet_ref.cpu() # 7.2 Optionally get Guidance Scale Embedding timestep_cond = None @@ -304,14 +316,14 @@ class HMImagePipeline(StableDiffusionImg2ImgPipeline): self.unet.cpu() - self.vae.to(device=device) + self.vae_decode.to(device=device) if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ + image = self.vae_decode.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[ 0 ] else: image = latents - self.vae.cpu() + self.vae_decode.cpu() do_denormalize = [True] * image.shape[0] diff --git a/hellomeme/pipelines/pipline_hm_video.py b/hellomeme/pipelines/pipline_hm_video.py index 9768969..d6dc9f1 100644 --- a/hellomeme/pipelines/pipline_hm_video.py +++ b/hellomeme/pipelines/pipline_hm_video.py @@ -63,7 +63,6 @@ class HMVideoPipeline(StableDiffusionImg2ImgPipeline): self.unet_pre.to(device='cpu', dtype=dtype).eval() if hasattr(self, "unet_ref"): - self.unet_ref.insert_reference_adapter(hm_adapter) self.unet_ref.to(device='cpu', dtype=dtype).eval() if hasattr(self, "mp_control"): diff --git a/meme.py b/meme.py index d08f62f..3903c39 100644 --- a/meme.py +++ b/meme.py @@ -57,7 +57,7 @@ def get_models_files(): return ['SD1.5'] + checkpoint_files, ['same as checkpoint', 'SD1.5 default vae'] + vae_files, ['None'] + lora_files -def append_pipline_weights(pipeline, checkpoint=None, lora=None, vae=None): +def append_pipline_weights(pipeline, checkpoint=None, lora=None, vae=None, stylize='x1'): ### load customized checkpoint or lora here: ## checkpoints @@ -72,13 +72,14 @@ def append_pipline_weights(pipeline, checkpoint=None, lora=None, vae=None): raw_stats = torch.load(checkpoint_path) if raw_stats: - state_dict = convert_ldm_unet_checkpoint(raw_stats, - pipeline.unet_pre.config if hasattr(pipeline, 'unet_pre') - else pipeline.unet.config) + state_dict = convert_ldm_unet_checkpoint(raw_stats, pipeline.unet_ref.config) if hasattr(pipeline, 'unet_pre'): pipeline.unet_pre.load_state_dict(state_dict, strict=False) pipeline.unet.load_state_dict(state_dict, strict=False) + if stylize == 'x2' and hasattr(pipeline, 'unet_ref'): + pipeline.unet_ref.load_state_dict(state_dict, strict=False) + if vae and not vae.startswith('SD1.5 default vae'): raw_vae_stats = raw_stats real_vae_path = '' @@ -97,7 +98,7 @@ def append_pipline_weights(pipeline, checkpoint=None, lora=None, vae=None): vae_state_dict = convert_ldm_vae_checkpoint(raw_vae_stats, pipeline.vae.config) if hasattr(pipeline, 'vae_decode'): pipeline.vae_decode.load_state_dict(vae_state_dict, strict=True) - else: + if stylize == 'x2': pipeline.vae.load_state_dict(vae_state_dict, strict=True) ### lora @@ -118,19 +119,20 @@ class HMImagePipelineLoader: "lora": (lora_files, ), "vae": (vae_files, ), "version": (['v1', 'v2'], ), + "stylize": (['x1', 'x2'], ), } } RETURN_TYPES = ("HMIMAGEPIPELINE", ) RETURN_NAMES = ("hm_image_pipeline", ) FUNCTION = "load_pipeline" CATEGORY = "hellomeme" - def load_pipeline(self, checkpoint=None, lora=None, vae=None, version='v2'): + def load_pipeline(self, checkpoint=None, lora=None, vae=None, version='v2', stylize='x1'): dtype = torch.float16 pipeline = HMImagePipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") pipeline.to(dtype=dtype) pipeline.caryomitosis(version=version) - append_pipline_weights(pipeline, checkpoint=checkpoint, lora=lora, vae=vae) + append_pipline_weights(pipeline, checkpoint=checkpoint, lora=lora, vae=vae, stylize=stylize) pipeline.insert_hm_modules(version=version, dtype=dtype) @@ -148,6 +150,7 @@ class HMVideoPipelineLoader: "lora": (lora_files, ), "vae": (vae_files, ), "version": (['v1', 'v2'], ), + "stylize": (['x1', 'x2'], ), } } @@ -156,13 +159,13 @@ class HMVideoPipelineLoader: FUNCTION = "load_pipeline" CATEGORY = "hellomeme" - def load_pipeline(self, checkpoint=None, lora=None, vae=None, version='v2'): + def load_pipeline(self, checkpoint=None, lora=None, vae=None, version='v2', stylize='x1'): dtype = torch.float16 pipeline = HMVideoPipeline.from_pretrained("stable-diffusion-v1-5/stable-diffusion-v1-5") pipeline.to(dtype=dtype) pipeline.caryomitosis(version=version) - append_pipline_weights(pipeline, checkpoint=checkpoint, lora=lora, vae=vae) + append_pipline_weights(pipeline, checkpoint=checkpoint, lora=lora, vae=vae, stylize=stylize) pipeline.insert_hm_modules(version=version, dtype=dtype)