Skip to content

GitLab

"tests/git@developer.sourcefind.cn:OpenDAS/diffusers.git" did not exist on "e55687e1e15407f60f32242027b7bb8170e58266"
Unverified Commit bdeff4d6 authored by cmdr2's avatar cmdr2 Committed by GitHub
Browse files

[ckpt loader] Allow loading the Inpaint and Img2Img pipelines, while loading a ckpt model (#2705) 

* [ckpt loader] Allow loading the Inpaint and Img2Img pipelines, while loading a ckpt model

* Address review comment from PR

* PyLint formatting

* Some more pylint fixes, unrelated to our change

* Another pylint fix

* Styling fix
parent fc188391
Hide whitespace changes
Inline Side-by-side
Showing with 78 additions and 19 deletions
src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py
View file @ bdeff4d6
...@@ -45,6 +45,8 @@ from diffusers import ( ...@@ -45,6 +45,8 @@ from diffusers import (
PNDMScheduler, PNDMScheduler,
PriorTransformer, PriorTransformer,
StableDiffusionControlNetPipeline, StableDiffusionControlNetPipeline,
StableDiffusionImg2ImgPipeline,
StableDiffusionInpaintPipeline,
StableDiffusionPipeline, StableDiffusionPipeline,
StableUnCLIPImg2ImgPipeline, StableUnCLIPImg2ImgPipeline,
StableUnCLIPPipeline, StableUnCLIPPipeline,
...@@ -979,6 +981,7 @@ def download_from_original_stable_diffusion_ckpt( ...@@ -979,6 +981,7 @@ def download_from_original_stable_diffusion_ckpt(
image_size: int = 512, image_size: int = 512,
prediction_type: str = None, prediction_type: str = None,
model_type: str = None, model_type: str = None,
is_img2img: bool = False,
extract_ema: bool = False, extract_ema: bool = False,
scheduler_type: str = "pndm", scheduler_type: str = "pndm",
num_in_channels: Optional[int] = None, num_in_channels: Optional[int] = None,
...@@ -1018,6 +1021,8 @@ def download_from_original_stable_diffusion_ckpt( ...@@ -1018,6 +1021,8 @@ def download_from_original_stable_diffusion_ckpt(
model_type (`str`, *optional*, defaults to `None`): model_type (`str`, *optional*, defaults to `None`):
The pipeline type. `None` to automatically infer, or one of `["FrozenOpenCLIPEmbedder", The pipeline type. `None` to automatically infer, or one of `["FrozenOpenCLIPEmbedder",
"FrozenCLIPEmbedder", "PaintByExample"]`. "FrozenCLIPEmbedder", "PaintByExample"]`.
is_img2img (`bool`, *optional*, defaults to `False`):
Whether the model should be loaded as an img2img pipeline.
extract_ema (`bool`, *optional*, defaults to `False`): Only relevant for extract_ema (`bool`, *optional*, defaults to `False`): Only relevant for
checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights or not. Defaults to checkpoints that have both EMA and non-EMA weights. Whether to extract the EMA weights or not. Defaults to
`False`. Pass `True` to extract the EMA weights. EMA weights usually yield higher quality images for `False`. Pass `True` to extract the EMA weights. EMA weights usually yield higher quality images for
...@@ -1193,16 +1198,44 @@ def download_from_original_stable_diffusion_ckpt( ...@@ -1193,16 +1198,44 @@ def download_from_original_stable_diffusion_ckpt(
requires_safety_checker=False, requires_safety_checker=False,
) )
else: else:
pipe = StableDiffusionPipeline( if (
vae=vae, hasattr(original_config, "model")
text_encoder=text_model, and hasattr(original_config.model, "target")
tokenizer=tokenizer, and "LatentInpaintDiffusion" in original_config.model.target
unet=unet, ):
scheduler=scheduler, pipe = StableDiffusionInpaintPipeline(
safety_checker=None, vae=vae,
feature_extractor=None, text_encoder=text_model,
requires_safety_checker=False, tokenizer=tokenizer,
) unet=unet,
scheduler=scheduler,
safety_checker=None,
feature_extractor=None,
requires_safety_checker=False,
)
else:
if is_img2img:
pipe = StableDiffusionImg2ImgPipeline(
vae=vae,
text_encoder=text_model,
tokenizer=tokenizer,
unet=unet,
scheduler=scheduler,
safety_checker=None,
feature_extractor=None,
requires_safety_checker=False,
)
else:
pipe = StableDiffusionPipeline(
vae=vae,
text_encoder=text_model,
tokenizer=tokenizer,
unet=unet,
scheduler=scheduler,
safety_checker=None,
feature_extractor=None,
requires_safety_checker=False,
)
else: else:
image_normalizer, image_noising_scheduler = stable_unclip_image_noising_components( image_normalizer, image_noising_scheduler = stable_unclip_image_noising_components(
original_config, clip_stats_path=clip_stats_path, device=device original_config, clip_stats_path=clip_stats_path, device=device
...@@ -1293,15 +1326,41 @@ def download_from_original_stable_diffusion_ckpt( ...@@ -1293,15 +1326,41 @@ def download_from_original_stable_diffusion_ckpt(
feature_extractor=feature_extractor, feature_extractor=feature_extractor,
) )
else: else:
pipe = StableDiffusionPipeline( if (
vae=vae, hasattr(original_config, "model")
text_encoder=text_model, and hasattr(original_config.model, "target")
tokenizer=tokenizer, and "LatentInpaintDiffusion" in original_config.model.target
unet=unet, ):
scheduler=scheduler, pipe = StableDiffusionInpaintPipeline(
safety_checker=safety_checker, vae=vae,
feature_extractor=feature_extractor, text_encoder=text_model,
) tokenizer=tokenizer,
unet=unet,
scheduler=scheduler,
safety_checker=safety_checker,
feature_extractor=feature_extractor,
)
else:
if is_img2img:
pipe = StableDiffusionImg2ImgPipeline(
vae=vae,
text_encoder=text_model,
tokenizer=tokenizer,
unet=unet,
scheduler=scheduler,
safety_checker=safety_checker,
feature_extractor=feature_extractor,
)
else:
pipe = StableDiffusionPipeline(
vae=vae,
text_encoder=text_model,
tokenizer=tokenizer,
unet=unet,
scheduler=scheduler,
safety_checker=safety_checker,
feature_extractor=feature_extractor,
)
else: else:
text_config = create_ldm_bert_config(original_config) text_config = create_ldm_bert_config(original_config)
text_model = convert_ldm_bert_checkpoint(checkpoint, text_config) text_model = convert_ldm_bert_checkpoint(checkpoint, text_config)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment