Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
patrickvonplaten merged 7 commits into from
Apr 19, 2023
Merged

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

Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
c53fc3f
[ckpt loader] Allow loading the Inpaint and Img2Img pipelines, while …
cmdr2 Mar 16, 2023
d3d6c32
Address review comment from PR
cmdr2 Mar 23, 2023
11accc7
PyLint formatting
cmdr2 Mar 30, 2023
6ab1317
Some more pylint fixes, unrelated to our change
cmdr2 Mar 30, 2023
9702400
Another pylint fix
cmdr2 Mar 30, 2023
04e7109
Merge pull request #2 from huggingface/main
cmdr2 Mar 30, 2023
464831b
Styling fix
cmdr2 Mar 30, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 78 additions & 19 deletions src/diffusers/pipelines/stable_diffusion/convert_from_ckpt.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,8 @@
PNDMScheduler,
PriorTransformer,
StableDiffusionControlNetPipeline,
StableDiffusionImg2ImgPipeline,
StableDiffusionInpaintPipeline,
StableDiffusionPipeline,
StableUnCLIPImg2ImgPipeline,
StableUnCLIPPipeline,
Expand Down Expand Up @@ -979,6 +981,7 @@ def download_from_original_stable_diffusion_ckpt(
image_size: int = 512,
prediction_type: str = None,
model_type: str = None,
is_img2img: bool = False,
Copy link
Contributor

@patrickvonplaten patrickvonplaten Mar 16, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ok for me!

extract_ema: bool = False,
scheduler_type: str = "pndm",
num_in_channels: Optional[int] = None,
Expand Down Expand Up @@ -1018,6 +1021,8 @@ def download_from_original_stable_diffusion_ckpt(
model_type (`str`, *optional*, defaults to `None`):
The pipeline type. `None` to automatically infer, or one of `["FrozenOpenCLIPEmbedder",
"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
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
Expand Down Expand Up @@ -1193,16 +1198,44 @@ def download_from_original_stable_diffusion_ckpt(
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,
)
if (
hasattr(original_config, "model")
and hasattr(original_config.model, "target")
and "LatentInpaintDiffusion" in original_config.model.target
):
pipe = StableDiffusionInpaintPipeline(
vae=vae,
text_encoder=text_model,
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:
image_normalizer, image_noising_scheduler = stable_unclip_image_noising_components(
original_config, clip_stats_path=clip_stats_path, device=device
Expand Down Expand Up @@ -1293,15 +1326,41 @@ def download_from_original_stable_diffusion_ckpt(
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,
)
if (
hasattr(original_config, "model")
and hasattr(original_config.model, "target")
and "LatentInpaintDiffusion" in original_config.model.target
):
pipe = StableDiffusionInpaintPipeline(
vae=vae,
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:
text_config = create_ldm_bert_config(original_config)
text_model = convert_ldm_bert_checkpoint(checkpoint, text_config)
Expand Down