nodes/HelloVision-ComfyUI_HelloMeme
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

add stylize option for pipeline loader

Browse Source
This commit is contained in:
songkey-win
2024-12-10 17:49:41 +08:00
parent 61a07d5a28
commit 87b3264bf5
3 changed files with 29 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
hellomeme/pipelines/pipline_hm_image.py
+17 -5
View File
@@ -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]
hellomeme/pipelines/pipline_hm_video.py
-1
View File
@@ -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"):
meme.py
+12 -9
View File
@@ -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)