[LCM] Make sure img2img works (huggingface#5632)

* [LCM] Clean up implementations * Add all * correct more * correct more * finish * up
kashif · Nov 2, 2023 · 072e008 · 072e008
1 parent b91d5dd
commit 072e008
Show file tree

Hide file tree

Showing 10 changed files with 972 additions and 6 deletions.
diff --git a/docs/source/en/api/pipelines/latent_consistency_models.md b/docs/source/en/api/pipelines/latent_consistency_models.md
@@ -10,6 +10,8 @@ A demo for the [SimianLuo/LCM_Dreamshaper_v7](https://huggingface.co/SimianLuo/L
 
 This pipeline was contributed by [luosiallen](https://luosiallen.github.io/) and [dg845](https://github.com/dg845).
 
+## text-to-image
+
 ```python
 import torch
 from diffusers import DiffusionPipeline
@@ -27,6 +29,27 @@ num_inference_steps = 4
 images = pipe(prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=8.0).images
 ```
 
+## image-to-image
+
+```python
+import torch
+from diffusers import AutoPipelineForImage2Image
+import PIL
+
+pipe = AutoPipelineForImage2Image.from_pretrained("SimianLuo/LCM_Dreamshaper_v7", torch_dtype=torch.float32)
+
+# To save GPU memory, torch.float16 can be used, but it may compromise image quality.
+pipe.to(torch_device="cuda", torch_dtype=torch.float32)
+
+prompt = "High altitude snowy mountains"
+image = PIL.Image.open("./snowy_mountains.png")
+
+# Can be set to 1~50 steps. LCM support fast inference even <= 4 steps. Recommend: 1~8 steps.
+num_inference_steps = 4 
+
+images = pipe(prompt=prompt, image=image, num_inference_steps=num_inference_steps, guidance_scale=8.0).images
+```
+
 ## LatentConsistencyModelPipeline
 
 [[autodoc]] LatentConsistencyModelPipeline
@@ -39,6 +62,16 @@ images = pipe(prompt=prompt, num_inference_steps=num_inference_steps, guidance_s
     - enable_vae_tiling
     - disable_vae_tiling
 
+[[autodoc]] LatentConsistencyModelImg2ImgPipeline
+    - all
+    - __call__
+    - enable_freeu
+    - disable_freeu
+    - enable_vae_slicing
+    - disable_vae_slicing
+    - enable_vae_tiling
+    - disable_vae_tiling
+
 ## StableDiffusionPipelineOutput
 
 [[autodoc]] pipelines.stable_diffusion.StableDiffusionPipelineOutput
diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py
@@ -230,6 +230,7 @@
             "KandinskyV22Pipeline",
             "KandinskyV22PriorEmb2EmbPipeline",
             "KandinskyV22PriorPipeline",
+            "LatentConsistencyModelImg2ImgPipeline",
             "LatentConsistencyModelPipeline",
             "LDMTextToImagePipeline",
             "MusicLDMPipeline",
@@ -573,6 +574,7 @@
             KandinskyV22Pipeline,
             KandinskyV22PriorEmb2EmbPipeline,
             KandinskyV22PriorPipeline,
+            LatentConsistencyModelImg2ImgPipeline,
             LatentConsistencyModelPipeline,
             LDMTextToImagePipeline,
             MusicLDMPipeline,

diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py
@@ -110,7 +110,10 @@
         "KandinskyV22PriorEmb2EmbPipeline",
         "KandinskyV22PriorPipeline",
     ]
-    _import_structure["latent_consistency_models"] = ["LatentConsistencyModelPipeline"]
+    _import_structure["latent_consistency_models"] = [
+        "LatentConsistencyModelImg2ImgPipeline",
+        "LatentConsistencyModelPipeline",
+    ]
     _import_structure["latent_diffusion"].extend(["LDMTextToImagePipeline"])
     _import_structure["musicldm"] = ["MusicLDMPipeline"]
     _import_structure["paint_by_example"] = ["PaintByExamplePipeline"]
@@ -334,7 +337,7 @@
             KandinskyV22PriorEmb2EmbPipeline,
             KandinskyV22PriorPipeline,
         )
-        from .latent_consistency_models import LatentConsistencyModelPipeline
+        from .latent_consistency_models import LatentConsistencyModelImg2ImgPipeline, LatentConsistencyModelPipeline
         from .latent_diffusion import LDMTextToImagePipeline
         from .musicldm import MusicLDMPipeline
         from .paint_by_example import PaintByExamplePipeline

diff --git a/src/diffusers/pipelines/auto_pipeline.py b/src/diffusers/pipelines/auto_pipeline.py
@@ -42,6 +42,7 @@
     KandinskyV22InpaintPipeline,
     KandinskyV22Pipeline,
 )
+from .latent_consistency_models import LatentConsistencyModelImg2ImgPipeline, LatentConsistencyModelPipeline
 from .stable_diffusion import (
     StableDiffusionImg2ImgPipeline,
     StableDiffusionInpaintPipeline,
@@ -65,6 +66,7 @@
         ("stable-diffusion-controlnet", StableDiffusionControlNetPipeline),
         ("stable-diffusion-xl-controlnet", StableDiffusionXLControlNetPipeline),
         ("wuerstchen", WuerstchenCombinedPipeline),
+        ("lcm", LatentConsistencyModelPipeline),
     ]
 )
 
@@ -77,6 +79,7 @@
         ("kandinsky22", KandinskyV22Img2ImgCombinedPipeline),
         ("stable-diffusion-controlnet", StableDiffusionControlNetImg2ImgPipeline),
         ("stable-diffusion-xl-controlnet", StableDiffusionXLControlNetImg2ImgPipeline),
+        ("lcm", LatentConsistencyModelImg2ImgPipeline),
     ]
 )
 

diff --git a/src/diffusers/pipelines/latent_consistency_models/__init__.py b/src/diffusers/pipelines/latent_consistency_models/__init__.py
@@ -5,11 +5,15 @@
 )
 
 
-_import_structure = {"pipeline_latent_consistency_models": ["LatentConsistencyModelPipeline"]}
+_import_structure = {
+    "pipeline_latent_consistency_img2img": ["LatentConsistencyModelImg2ImgPipeline"],
+    "pipeline_latent_consistency_text2img": ["LatentConsistencyModelPipeline"],
+}
 
 
 if TYPE_CHECKING:
-    from .pipeline_latent_consistency_models import LatentConsistencyModelPipeline
+    from .pipeline_latent_consistency_img2img import LatentConsistencyModelImg2ImgPipeline
+    from .pipeline_latent_consistency_text2img import LatentConsistencyModelPipeline
 
 else:
     import sys