Spaces:
Running
Running
Upload space
Browse files- LICENSE +21 -0
- README.md +9 -6
- app.py +176 -0
- lcm_ov_pipeline.py +388 -0
- lcm_scheduler.py +529 -0
- requirements.txt +12 -0
- style.css +16 -0
LICENSE
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
MIT License
|
2 |
+
|
3 |
+
Copyright (c) 2023 hysts
|
4 |
+
|
5 |
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
6 |
+
of this software and associated documentation files (the "Software"), to deal
|
7 |
+
in the Software without restriction, including without limitation the rights
|
8 |
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
9 |
+
copies of the Software, and to permit persons to whom the Software is
|
10 |
+
furnished to do so, subject to the following conditions:
|
11 |
+
|
12 |
+
The above copyright notice and this permission notice shall be included in all
|
13 |
+
copies or substantial portions of the Software.
|
14 |
+
|
15 |
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
16 |
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
17 |
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
18 |
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
19 |
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
20 |
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
21 |
+
SOFTWARE.
|
README.md
CHANGED
@@ -1,13 +1,16 @@
|
|
1 |
---
|
2 |
-
title: Latent Consistency
|
3 |
-
emoji:
|
4 |
-
colorFrom:
|
5 |
-
colorTo:
|
6 |
sdk: gradio
|
7 |
-
sdk_version: 3.
|
8 |
app_file: app.py
|
9 |
-
pinned: false
|
10 |
license: mit
|
|
|
|
|
|
|
|
|
11 |
---
|
12 |
|
13 |
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
|
|
1 |
---
|
2 |
+
title: Latent Consistency Models OpenVino CPU
|
3 |
+
emoji: 🥶
|
4 |
+
colorFrom: blue
|
5 |
+
colorTo: purple
|
6 |
sdk: gradio
|
7 |
+
sdk_version: 3.48.0
|
8 |
app_file: app.py
|
|
|
9 |
license: mit
|
10 |
+
pinned: true
|
11 |
+
suggested_hardware: cpu-basic
|
12 |
+
suggested_storage: small
|
13 |
+
hf_oauth: true
|
14 |
---
|
15 |
|
16 |
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
app.py
ADDED
@@ -0,0 +1,176 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#!/usr/bin/env python
|
2 |
+
from __future__ import annotations
|
3 |
+
|
4 |
+
import os
|
5 |
+
import random
|
6 |
+
import time
|
7 |
+
|
8 |
+
import gradio as gr
|
9 |
+
import numpy as np
|
10 |
+
import PIL.Image
|
11 |
+
|
12 |
+
from diffusers import DiffusionPipeline
|
13 |
+
|
14 |
+
from lcm_scheduler import LCMScheduler
|
15 |
+
from lcm_ov_pipeline import OVLatentConsistencyModelPipeline
|
16 |
+
|
17 |
+
import os
|
18 |
+
from tqdm import tqdm
|
19 |
+
import gradio_user_history as gr_user_history
|
20 |
+
|
21 |
+
from concurrent.futures import ThreadPoolExecutor
|
22 |
+
import uuid
|
23 |
+
|
24 |
+
DESCRIPTION = '''# Latent Consistency Model OpenVino CPU
|
25 |
+
Based on [Latency Consistency Model](https://huggingface.co/spaces/SimianLuo/Latent_Consistency_Model) HF space
|
26 |
+
|
27 |
+
Distilled from [Dreamshaper v7](https://huggingface.co/Lykon/dreamshaper-7) fine-tune of [Stable Diffusion v1-5](https://huggingface.co/runwayml/stable-diffusion-v1-5) with only 4,000 training iterations (~32 A100 GPU Hours). [Project page](https://latent-consistency-models.github.io)
|
28 |
+
|
29 |
+
<p>Running on CPU 🥶.</p>
|
30 |
+
'''
|
31 |
+
|
32 |
+
MAX_SEED = np.iinfo(np.int32).max
|
33 |
+
|
34 |
+
model_id = "deinferno/LCM_Dreamshaper_v7-openvino"
|
35 |
+
batch_size = 1
|
36 |
+
width = int(os.getenv("IMAGE_WIDTH", "512"))
|
37 |
+
height = int(os.getenv("IMAGE_HEIGHT", "512"))
|
38 |
+
num_images = int(os.getenv("NUM_IMAGES", "1"))
|
39 |
+
|
40 |
+
scheduler = LCMScheduler.from_pretrained(model_id, subfolder="scheduler")
|
41 |
+
pipe = OVLatentConsistencyModelPipeline.from_pretrained(model_id, scheduler = scheduler, compile = False)
|
42 |
+
pipe.reshape(batch_size=batch_size, height=height, width=width, num_images_per_prompt=num_images)
|
43 |
+
pipe.compile()
|
44 |
+
|
45 |
+
def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
|
46 |
+
if randomize_seed:
|
47 |
+
seed = random.randint(0, MAX_SEED)
|
48 |
+
return seed
|
49 |
+
|
50 |
+
def save_image(img, profile: gr.OAuthProfile | None, metadata: dict):
|
51 |
+
unique_name = str(uuid.uuid4()) + '.png'
|
52 |
+
img.save(unique_name)
|
53 |
+
gr_user_history.save_image(label=metadata["prompt"], image=img, profile=profile, metadata=metadata)
|
54 |
+
return unique_name
|
55 |
+
|
56 |
+
def save_images(image_array, profile: gr.OAuthProfile | None, metadata: dict):
|
57 |
+
paths = []
|
58 |
+
with ThreadPoolExecutor() as executor:
|
59 |
+
paths = list(executor.map(save_image, image_array, [profile]*len(image_array), [metadata]*len(image_array)))
|
60 |
+
return paths
|
61 |
+
|
62 |
+
def generate(
|
63 |
+
prompt: str,
|
64 |
+
seed: int = 0,
|
65 |
+
guidance_scale: float = 8.0,
|
66 |
+
num_inference_steps: int = 4,
|
67 |
+
randomize_seed: bool = False,
|
68 |
+
progress = gr.Progress(track_tqdm=True),
|
69 |
+
profile: gr.OAuthProfile | None = None,
|
70 |
+
) -> PIL.Image.Image:
|
71 |
+
global batch_size
|
72 |
+
global width
|
73 |
+
global height
|
74 |
+
global num_images
|
75 |
+
|
76 |
+
seed = randomize_seed_fn(seed, randomize_seed)
|
77 |
+
np.random.seed(seed)
|
78 |
+
start_time = time.time()
|
79 |
+
result = pipe(
|
80 |
+
prompt=prompt,
|
81 |
+
width=width,
|
82 |
+
height=height,
|
83 |
+
guidance_scale=guidance_scale,
|
84 |
+
num_inference_steps=num_inference_steps,
|
85 |
+
num_images_per_prompt=num_images,
|
86 |
+
output_type="pil",
|
87 |
+
).images
|
88 |
+
paths = save_images(result, profile, metadata={"prompt": prompt, "seed": seed, "width": width, "height": height, "guidance_scale": guidance_scale, "num_inference_steps": num_inference_steps})
|
89 |
+
print(time.time() - start_time)
|
90 |
+
return paths, seed
|
91 |
+
|
92 |
+
examples = [
|
93 |
+
"portrait photo of a girl, photograph, highly detailed face, depth of field, moody light, golden hour, style by Dan Winters, Russell James, Steve McCurry, centered, extremely detailed, Nikon D850, award winning photography",
|
94 |
+
"Self-portrait oil painting, a beautiful cyborg with golden hair, 8k",
|
95 |
+
"Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
|
96 |
+
"A photo of beautiful mountain with realistic sunset and blue lake, highly detailed, masterpiece",
|
97 |
+
]
|
98 |
+
|
99 |
+
with gr.Blocks(css="style.css") as demo:
|
100 |
+
gr.Markdown(DESCRIPTION)
|
101 |
+
gr.DuplicateButton(
|
102 |
+
value="Duplicate Space for private use",
|
103 |
+
elem_id="duplicate-button",
|
104 |
+
visible=os.getenv("SHOW_DUPLICATE_BUTTON") == "1",
|
105 |
+
)
|
106 |
+
with gr.Group():
|
107 |
+
with gr.Row():
|
108 |
+
prompt = gr.Text(
|
109 |
+
label="Prompt",
|
110 |
+
show_label=False,
|
111 |
+
max_lines=1,
|
112 |
+
placeholder="Enter your prompt",
|
113 |
+
container=False,
|
114 |
+
)
|
115 |
+
run_button = gr.Button("Run", scale=0)
|
116 |
+
result = gr.Gallery(
|
117 |
+
label="Generated images", show_label=False, elem_id="gallery", grid=[2]
|
118 |
+
)
|
119 |
+
with gr.Accordion("Advanced options", open=False):
|
120 |
+
seed = gr.Slider(
|
121 |
+
label="Seed",
|
122 |
+
minimum=0,
|
123 |
+
maximum=MAX_SEED,
|
124 |
+
step=1,
|
125 |
+
value=0,
|
126 |
+
randomize=True
|
127 |
+
)
|
128 |
+
randomize_seed = gr.Checkbox(label="Randomize seed across runs", value=True)
|
129 |
+
with gr.Row():
|
130 |
+
guidance_scale = gr.Slider(
|
131 |
+
label="Guidance scale for base",
|
132 |
+
minimum=2,
|
133 |
+
maximum=14,
|
134 |
+
step=0.1,
|
135 |
+
value=8.0,
|
136 |
+
)
|
137 |
+
num_inference_steps = gr.Slider(
|
138 |
+
label="Number of inference steps for base",
|
139 |
+
minimum=1,
|
140 |
+
maximum=8,
|
141 |
+
step=1,
|
142 |
+
value=4,
|
143 |
+
)
|
144 |
+
|
145 |
+
with gr.Accordion("Past generations", open=False):
|
146 |
+
gr_user_history.render()
|
147 |
+
|
148 |
+
gr.Examples(
|
149 |
+
examples=examples,
|
150 |
+
inputs=prompt,
|
151 |
+
outputs=result,
|
152 |
+
fn=generate,
|
153 |
+
cache_examples=CACHE_EXAMPLES,
|
154 |
+
)
|
155 |
+
|
156 |
+
gr.on(
|
157 |
+
triggers=[
|
158 |
+
prompt.submit,
|
159 |
+
run_button.click,
|
160 |
+
],
|
161 |
+
fn=generate,
|
162 |
+
inputs=[
|
163 |
+
prompt,
|
164 |
+
seed,
|
165 |
+
guidance_scale,
|
166 |
+
num_inference_steps,
|
167 |
+
randomize_seed
|
168 |
+
],
|
169 |
+
outputs=[result, seed],
|
170 |
+
api_name="run",
|
171 |
+
)
|
172 |
+
|
173 |
+
if __name__ == "__main__":
|
174 |
+
demo.queue(api_open=False)
|
175 |
+
# demo.queue(max_size=20).launch()
|
176 |
+
demo.launch()
|
lcm_ov_pipeline.py
ADDED
@@ -0,0 +1,388 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import inspect
|
2 |
+
|
3 |
+
from pathlib import Path
|
4 |
+
from tempfile import TemporaryDirectory
|
5 |
+
from typing import List, Optional, Tuple, Union, Dict, Any, Callable, OrderedDict
|
6 |
+
|
7 |
+
import numpy as np
|
8 |
+
import openvino
|
9 |
+
import torch
|
10 |
+
|
11 |
+
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
|
12 |
+
from optimum.intel.openvino.modeling_diffusion import OVStableDiffusionPipeline, OVModelUnet, OVModelVaeDecoder, OVModelTextEncoder, OVModelVaeEncoder, VaeImageProcessor
|
13 |
+
from optimum.utils import (
|
14 |
+
DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER,
|
15 |
+
DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER,
|
16 |
+
DIFFUSION_MODEL_UNET_SUBFOLDER,
|
17 |
+
DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER,
|
18 |
+
DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER,
|
19 |
+
)
|
20 |
+
|
21 |
+
|
22 |
+
from diffusers import logging
|
23 |
+
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
24 |
+
|
25 |
+
class LCMOVModelUnet(OVModelUnet):
|
26 |
+
def __call__(
|
27 |
+
self,
|
28 |
+
sample: np.ndarray,
|
29 |
+
timestep: np.ndarray,
|
30 |
+
encoder_hidden_states: np.ndarray,
|
31 |
+
timestep_cond: Optional[np.ndarray] = None,
|
32 |
+
text_embeds: Optional[np.ndarray] = None,
|
33 |
+
time_ids: Optional[np.ndarray] = None,
|
34 |
+
):
|
35 |
+
self._compile()
|
36 |
+
|
37 |
+
inputs = {
|
38 |
+
"sample": sample,
|
39 |
+
"timestep": timestep,
|
40 |
+
"encoder_hidden_states": encoder_hidden_states,
|
41 |
+
}
|
42 |
+
|
43 |
+
if timestep_cond is not None:
|
44 |
+
inputs["timestep_cond"] = timestep_cond
|
45 |
+
if text_embeds is not None:
|
46 |
+
inputs["text_embeds"] = text_embeds
|
47 |
+
if time_ids is not None:
|
48 |
+
inputs["time_ids"] = time_ids
|
49 |
+
|
50 |
+
outputs = self.request(inputs, shared_memory=True)
|
51 |
+
return list(outputs.values())
|
52 |
+
|
53 |
+
class OVLatentConsistencyModelPipeline(OVStableDiffusionPipeline):
|
54 |
+
|
55 |
+
def __init__(
|
56 |
+
self,
|
57 |
+
vae_decoder: openvino.runtime.Model,
|
58 |
+
text_encoder: openvino.runtime.Model,
|
59 |
+
unet: openvino.runtime.Model,
|
60 |
+
config: Dict[str, Any],
|
61 |
+
tokenizer: "CLIPTokenizer",
|
62 |
+
scheduler: Union["DDIMScheduler", "PNDMScheduler", "LMSDiscreteScheduler"],
|
63 |
+
feature_extractor: Optional["CLIPFeatureExtractor"] = None,
|
64 |
+
vae_encoder: Optional[openvino.runtime.Model] = None,
|
65 |
+
text_encoder_2: Optional[openvino.runtime.Model] = None,
|
66 |
+
tokenizer_2: Optional["CLIPTokenizer"] = None,
|
67 |
+
device: str = "CPU",
|
68 |
+
dynamic_shapes: bool = True,
|
69 |
+
compile: bool = True,
|
70 |
+
ov_config: Optional[Dict[str, str]] = None,
|
71 |
+
model_save_dir: Optional[Union[str, Path, TemporaryDirectory]] = None,
|
72 |
+
**kwargs,
|
73 |
+
):
|
74 |
+
self._internal_dict = config
|
75 |
+
self._device = device.upper()
|
76 |
+
self.is_dynamic = dynamic_shapes
|
77 |
+
self.ov_config = ov_config if ov_config is not None else {}
|
78 |
+
self._model_save_dir = (
|
79 |
+
Path(model_save_dir.name) if isinstance(model_save_dir, TemporaryDirectory) else model_save_dir
|
80 |
+
)
|
81 |
+
self.vae_decoder = OVModelVaeDecoder(vae_decoder, self)
|
82 |
+
self.unet = LCMOVModelUnet(unet, self)
|
83 |
+
self.text_encoder = OVModelTextEncoder(text_encoder, self) if text_encoder is not None else None
|
84 |
+
self.text_encoder_2 = (
|
85 |
+
OVModelTextEncoder(text_encoder_2, self, model_name=DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER)
|
86 |
+
if text_encoder_2 is not None
|
87 |
+
else None
|
88 |
+
)
|
89 |
+
self.vae_encoder = OVModelVaeEncoder(vae_encoder, self) if vae_encoder is not None else None
|
90 |
+
|
91 |
+
if "block_out_channels" in self.vae_decoder.config:
|
92 |
+
self.vae_scale_factor = 2 ** (len(self.vae_decoder.config["block_out_channels"]) - 1)
|
93 |
+
else:
|
94 |
+
self.vae_scale_factor = 8
|
95 |
+
|
96 |
+
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
|
97 |
+
|
98 |
+
self.tokenizer = tokenizer
|
99 |
+
self.tokenizer_2 = tokenizer_2
|
100 |
+
self.scheduler = scheduler
|
101 |
+
self.feature_extractor = feature_extractor
|
102 |
+
self.safety_checker = None
|
103 |
+
self.preprocessors = []
|
104 |
+
|
105 |
+
if self.is_dynamic:
|
106 |
+
self.reshape(batch_size=-1, height=-1, width=-1, num_images_per_prompt=-1)
|
107 |
+
|
108 |
+
if compile:
|
109 |
+
self.compile()
|
110 |
+
|
111 |
+
sub_models = {
|
112 |
+
DIFFUSION_MODEL_TEXT_ENCODER_SUBFOLDER: self.text_encoder,
|
113 |
+
DIFFUSION_MODEL_UNET_SUBFOLDER: self.unet,
|
114 |
+
DIFFUSION_MODEL_VAE_DECODER_SUBFOLDER: self.vae_decoder,
|
115 |
+
DIFFUSION_MODEL_VAE_ENCODER_SUBFOLDER: self.vae_encoder,
|
116 |
+
DIFFUSION_MODEL_TEXT_ENCODER_2_SUBFOLDER: self.text_encoder_2,
|
117 |
+
}
|
118 |
+
for name in sub_models.keys():
|
119 |
+
self._internal_dict[name] = (
|
120 |
+
("optimum", sub_models[name].__class__.__name__) if sub_models[name] is not None else (None, None)
|
121 |
+
)
|
122 |
+
|
123 |
+
self._internal_dict.pop("vae", None)
|
124 |
+
|
125 |
+
def _reshape_unet(
|
126 |
+
self,
|
127 |
+
model: openvino.runtime.Model,
|
128 |
+
batch_size: int = -1,
|
129 |
+
height: int = -1,
|
130 |
+
width: int = -1,
|
131 |
+
num_images_per_prompt: int = -1,
|
132 |
+
tokenizer_max_length: int = -1,
|
133 |
+
):
|
134 |
+
if batch_size == -1 or num_images_per_prompt == -1:
|
135 |
+
batch_size = -1
|
136 |
+
else:
|
137 |
+
batch_size = batch_size * num_images_per_prompt
|
138 |
+
|
139 |
+
height = height // self.vae_scale_factor if height > 0 else height
|
140 |
+
width = width // self.vae_scale_factor if width > 0 else width
|
141 |
+
shapes = {}
|
142 |
+
for inputs in model.inputs:
|
143 |
+
shapes[inputs] = inputs.get_partial_shape()
|
144 |
+
if inputs.get_any_name() == "timestep":
|
145 |
+
shapes[inputs][0] = 1
|
146 |
+
elif inputs.get_any_name() == "sample":
|
147 |
+
in_channels = self.unet.config.get("in_channels", None)
|
148 |
+
if in_channels is None:
|
149 |
+
in_channels = shapes[inputs][1]
|
150 |
+
if in_channels.is_dynamic:
|
151 |
+
logger.warning(
|
152 |
+
"Could not identify `in_channels` from the unet configuration, to statically reshape the unet please provide a configuration."
|
153 |
+
)
|
154 |
+
self.is_dynamic = True
|
155 |
+
|
156 |
+
shapes[inputs] = [batch_size, in_channels, height, width]
|
157 |
+
elif inputs.get_any_name() == "timestep_cond":
|
158 |
+
shapes[inputs] = [batch_size, inputs.get_partial_shape()[1]]
|
159 |
+
elif inputs.get_any_name() == "text_embeds":
|
160 |
+
shapes[inputs] = [batch_size, self.text_encoder_2.config["projection_dim"]]
|
161 |
+
elif inputs.get_any_name() == "time_ids":
|
162 |
+
shapes[inputs] = [batch_size, inputs.get_partial_shape()[1]]
|
163 |
+
else:
|
164 |
+
shapes[inputs][0] = batch_size
|
165 |
+
shapes[inputs][1] = tokenizer_max_length
|
166 |
+
model.reshape(shapes)
|
167 |
+
return model
|
168 |
+
|
169 |
+
def get_guidance_scale_embedding(self, w, embedding_dim=512, dtype=np.float32):
|
170 |
+
"""
|
171 |
+
see https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
|
172 |
+
Args:
|
173 |
+
timesteps: np.array: generate embedding vectors at these timesteps
|
174 |
+
embedding_dim: int: dimension of the embeddings to generate
|
175 |
+
dtype: data type of the generated embeddings
|
176 |
+
|
177 |
+
Returns:
|
178 |
+
embedding vectors with shape `(len(timesteps), embedding_dim)`
|
179 |
+
"""
|
180 |
+
assert len(w.shape) == 1
|
181 |
+
w = w * 1000.
|
182 |
+
|
183 |
+
half_dim = embedding_dim // 2
|
184 |
+
emb = np.log(np.array(10000.)) / (half_dim - 1)
|
185 |
+
emb = np.exp(np.arange(half_dim, dtype=dtype) * -emb)
|
186 |
+
emb = w.astype(dtype)[:, None] * emb[None, :]
|
187 |
+
emb = np.concatenate([np.sin(emb), np.cos(emb)], axis=1)
|
188 |
+
if embedding_dim % 2 == 1: # zero pad
|
189 |
+
emb = np.pad(emb, (0, 1))
|
190 |
+
assert emb.shape == (w.shape[0], embedding_dim)
|
191 |
+
return emb
|
192 |
+
|
193 |
+
# Adapted from https://github.com/huggingface/optimum/blob/15b8d1eed4d83c5004d3b60f6b6f13744b358f01/optimum/pipelines/diffusers/pipeline_stable_diffusion.py#L201
|
194 |
+
def __call__(
|
195 |
+
self,
|
196 |
+
prompt: Optional[Union[str, List[str]]] = None,
|
197 |
+
height: Optional[int] = None,
|
198 |
+
width: Optional[int] = None,
|
199 |
+
num_inference_steps: int = 4,
|
200 |
+
original_inference_steps: int = None,
|
201 |
+
guidance_scale: float = 7.5,
|
202 |
+
num_images_per_prompt: int = 1,
|
203 |
+
eta: float = 0.0,
|
204 |
+
generator: Optional[np.random.RandomState] = None,
|
205 |
+
latents: Optional[np.ndarray] = None,
|
206 |
+
prompt_embeds: Optional[np.ndarray] = None,
|
207 |
+
output_type: str = "pil",
|
208 |
+
return_dict: bool = True,
|
209 |
+
callback: Optional[Callable[[int, int, np.ndarray], None]] = None,
|
210 |
+
callback_steps: int = 1,
|
211 |
+
guidance_rescale: float = 0.0,
|
212 |
+
):
|
213 |
+
r"""
|
214 |
+
Function invoked when calling the pipeline for generation.
|
215 |
+
|
216 |
+
Args:
|
217 |
+
prompt (`Optional[Union[str, List[str]]]`, defaults to None):
|
218 |
+
The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
|
219 |
+
instead.
|
220 |
+
height (`Optional[int]`, defaults to None):
|
221 |
+
The height in pixels of the generated image.
|
222 |
+
width (`Optional[int]`, defaults to None):
|
223 |
+
The width in pixels of the generated image.
|
224 |
+
num_inference_steps (`int`, defaults to 4):
|
225 |
+
The number of denoising steps. More denoising steps usually lead to a higher quality image at the
|
226 |
+
expense of slower inference.
|
227 |
+
original_inference_steps (`int`, *optional*):
|
228 |
+
The original number of inference steps use to generate a linearly-spaced timestep schedule, from which
|
229 |
+
we will draw `num_inference_steps` evenly spaced timesteps from as our final timestep schedule,
|
230 |
+
following the Skipping-Step method in the paper (see Section 4.3). If not set this will default to the
|
231 |
+
scheduler's `original_inference_steps` attribute.
|
232 |
+
guidance_scale (`float`, defaults to 7.5):
|
233 |
+
Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
|
234 |
+
`guidance_scale` is defined as `w` of equation 2. of [Imagen
|
235 |
+
Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
|
236 |
+
1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
|
237 |
+
usually at the expense of lower image quality.
|
238 |
+
num_images_per_prompt (`int`, defaults to 1):
|
239 |
+
The number of images to generate per prompt.
|
240 |
+
eta (`float`, defaults to 0.0):
|
241 |
+
Corresponds to parameter eta (η) in the DDIM paper: https://arxiv.org/abs/2010.02502. Only applies to
|
242 |
+
[`schedulers.DDIMScheduler`], will be ignored for others.
|
243 |
+
generator (`Optional[np.random.RandomState]`, defaults to `None`)::
|
244 |
+
A np.random.RandomState to make generation deterministic.
|
245 |
+
latents (`Optional[np.ndarray]`, defaults to `None`):
|
246 |
+
Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
|
247 |
+
generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
|
248 |
+
tensor will ge generated by sampling using the supplied random `generator`.
|
249 |
+
prompt_embeds (`Optional[np.ndarray]`, defaults to `None`):
|
250 |
+
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
|
251 |
+
provided, text embeddings will be generated from `prompt` input argument.
|
252 |
+
output_type (`str`, defaults to `"pil"`):
|
253 |
+
The output format of the generate image. Choose between
|
254 |
+
[PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
|
255 |
+
return_dict (`bool`, defaults to `True`):
|
256 |
+
Whether or not to return a [`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] instead of a
|
257 |
+
plain tuple.
|
258 |
+
callback (Optional[Callable], defaults to `None`):
|
259 |
+
A function that will be called every `callback_steps` steps during inference. The function will be
|
260 |
+
called with the following arguments: `callback(step: int, timestep: int, latents: torch.FloatTensor)`.
|
261 |
+
callback_steps (`int`, defaults to 1):
|
262 |
+
The frequency at which the `callback` function will be called. If not specified, the callback will be
|
263 |
+
called at every step.
|
264 |
+
guidance_rescale (`float`, defaults to 0.0):
|
265 |
+
Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are
|
266 |
+
Flawed](https://arxiv.org/pdf/2305.08891.pdf) `guidance_scale` is defined as `φ` in equation 16. of
|
267 |
+
[Common Diffusion Noise Schedules and Sample Steps are Flawed](https://arxiv.org/pdf/2305.08891.pdf).
|
268 |
+
Guidance rescale factor should fix overexposure when using zero terminal SNR.
|
269 |
+
|
270 |
+
Returns:
|
271 |
+
[`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] or `tuple`:
|
272 |
+
[`~pipelines.stable_diffusion.StableDiffusionPipelineOutput`] if `return_dict` is True, otherwise a `tuple.
|
273 |
+
When returning a tuple, the first element is a list with the generated images, and the second element is a
|
274 |
+
list of `bool`s denoting whether the corresponding generated image likely represents "not-safe-for-work"
|
275 |
+
(nsfw) content, according to the `safety_checker`.
|
276 |
+
"""
|
277 |
+
height = height or self.unet.config.get("sample_size", 64) * self.vae_scale_factor
|
278 |
+
width = width or self.unet.config.get("sample_size", 64) * self.vae_scale_factor
|
279 |
+
|
280 |
+
# check inputs. Raise error if not correct
|
281 |
+
self.check_inputs(
|
282 |
+
prompt, height, width, callback_steps, None, prompt_embeds, None
|
283 |
+
)
|
284 |
+
|
285 |
+
# define call parameters
|
286 |
+
if isinstance(prompt, str):
|
287 |
+
batch_size = 1
|
288 |
+
elif isinstance(prompt, list):
|
289 |
+
batch_size = len(prompt)
|
290 |
+
else:
|
291 |
+
batch_size = prompt_embeds.shape[0]
|
292 |
+
|
293 |
+
if generator is None:
|
294 |
+
generator = np.random
|
295 |
+
|
296 |
+
# Create torch.Generator instance with same state as np.random.RandomState
|
297 |
+
torch_generator = torch.Generator().manual_seed(int(generator.get_state()[1][0]))
|
298 |
+
|
299 |
+
#do_classifier_free_guidance = guidance_scale > 1.0
|
300 |
+
|
301 |
+
# NOTE: when a LCM is distilled from an LDM via latent consistency distillation (Algorithm 1) with guided
|
302 |
+
# distillation, the forward pass of the LCM learns to approximate sampling from the LDM using CFG with the
|
303 |
+
# unconditional prompt "" (the empty string). Due to this, LCMs currently do not support negative prompts.
|
304 |
+
prompt_embeds = self._encode_prompt(
|
305 |
+
prompt,
|
306 |
+
num_images_per_prompt,
|
307 |
+
False,
|
308 |
+
negative_prompt=None,
|
309 |
+
prompt_embeds=prompt_embeds,
|
310 |
+
negative_prompt_embeds=None,
|
311 |
+
)
|
312 |
+
|
313 |
+
# set timesteps
|
314 |
+
self.scheduler.set_timesteps(num_inference_steps, "cpu", original_inference_steps=original_inference_steps)
|
315 |
+
timesteps = self.scheduler.timesteps
|
316 |
+
|
317 |
+
latents = self.prepare_latents(
|
318 |
+
batch_size * num_images_per_prompt,
|
319 |
+
self.unet.config.get("in_channels", 4),
|
320 |
+
height,
|
321 |
+
width,
|
322 |
+
prompt_embeds.dtype,
|
323 |
+
generator,
|
324 |
+
latents,
|
325 |
+
)
|
326 |
+
|
327 |
+
# Get Guidance Scale Embedding
|
328 |
+
w = np.tile(guidance_scale - 1, batch_size * num_images_per_prompt)
|
329 |
+
w_embedding = self.get_guidance_scale_embedding(w, embedding_dim=self.unet.config.get("time_cond_proj_dim", 256))
|
330 |
+
|
331 |
+
# Adapted from diffusers to extend it for other runtimes than ORT
|
332 |
+
timestep_dtype = self.unet.input_dtype.get("timestep", np.float32)
|
333 |
+
|
334 |
+
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
|
335 |
+
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
|
336 |
+
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
|
337 |
+
# and should be between [0, 1]
|
338 |
+
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
|
339 |
+
extra_step_kwargs = {}
|
340 |
+
if accepts_eta:
|
341 |
+
extra_step_kwargs["eta"] = eta
|
342 |
+
|
343 |
+
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
|
344 |
+
if accepts_generator:
|
345 |
+
extra_step_kwargs["generator"] = torch_generator
|
346 |
+
|
347 |
+
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
|
348 |
+
for i, t in enumerate(self.progress_bar(timesteps)):
|
349 |
+
|
350 |
+
# predict the noise residual
|
351 |
+
timestep = np.array([t], dtype=timestep_dtype)
|
352 |
+
|
353 |
+
noise_pred = self.unet(sample=latents, timestep=timestep, timestep_cond = w_embedding, encoder_hidden_states=prompt_embeds)[0]
|
354 |
+
|
355 |
+
# compute the previous noisy sample x_t -> x_t-1
|
356 |
+
latents, denoised = self.scheduler.step(
|
357 |
+
torch.from_numpy(noise_pred), t, torch.from_numpy(latents), **extra_step_kwargs, return_dict = False
|
358 |
+
)
|
359 |
+
|
360 |
+
latents, denoised = latents.numpy(), denoised.numpy()
|
361 |
+
|
362 |
+
# call the callback, if provided
|
363 |
+
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
|
364 |
+
if callback is not None and i % callback_steps == 0:
|
365 |
+
callback(i, t, latents)
|
366 |
+
|
367 |
+
if output_type == "latent":
|
368 |
+
image = latents
|
369 |
+
has_nsfw_concept = None
|
370 |
+
else:
|
371 |
+
denoised /= self.vae_decoder.config.get("scaling_factor", 0.18215)
|
372 |
+
# it seems likes there is a strange result for using half-precision vae decoder if batchsize>1
|
373 |
+
image = np.concatenate(
|
374 |
+
[self.vae_decoder(latent_sample=denoised[i : i + 1])[0] for i in range(latents.shape[0])]
|
375 |
+
)
|
376 |
+
image, has_nsfw_concept = self.run_safety_checker(image)
|
377 |
+
|
378 |
+
if has_nsfw_concept is None:
|
379 |
+
do_denormalize = [True] * image.shape[0]
|
380 |
+
else:
|
381 |
+
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
|
382 |
+
|
383 |
+
image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
|
384 |
+
|
385 |
+
if not return_dict:
|
386 |
+
return (image, has_nsfw_concept)
|
387 |
+
|
388 |
+
return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
|
lcm_scheduler.py
ADDED
@@ -0,0 +1,529 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright 2023 Stanford University Team and The HuggingFace Team. All rights reserved.
|
2 |
+
#
|
3 |
+
# Licensed under the Apache License, Version 2.0 (the "License");
|
4 |
+
# you may not use this file except in compliance with the License.
|
5 |
+
# You may obtain a copy of the License at
|
6 |
+
#
|
7 |
+
# http://www.apache.org/licenses/LICENSE-2.0
|
8 |
+
#
|
9 |
+
# Unless required by applicable law or agreed to in writing, software
|
10 |
+
# distributed under the License is distributed on an "AS IS" BASIS,
|
11 |
+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
12 |
+
# See the License for the specific language governing permissions and
|
13 |
+
# limitations under the License.
|
14 |
+
|
15 |
+
# DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
|
16 |
+
# and https://github.com/hojonathanho/diffusion
|
17 |
+
|
18 |
+
import math
|
19 |
+
from dataclasses import dataclass
|
20 |
+
from typing import List, Optional, Tuple, Union
|
21 |
+
|
22 |
+
import numpy as np
|
23 |
+
import torch
|
24 |
+
|
25 |
+
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
26 |
+
from diffusers.utils import BaseOutput, logging
|
27 |
+
from diffusers.utils.torch_utils import randn_tensor
|
28 |
+
from diffusers.schedulers.scheduling_utils import SchedulerMixin
|
29 |
+
|
30 |
+
|
31 |
+
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
32 |
+
|
33 |
+
|
34 |
+
@dataclass
|
35 |
+
class LCMSchedulerOutput(BaseOutput):
|
36 |
+
"""
|
37 |
+
Output class for the scheduler's `step` function output.
|
38 |
+
|
39 |
+
Args:
|
40 |
+
prev_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
|
41 |
+
Computed sample `(x_{t-1})` of previous timestep. `prev_sample` should be used as next model input in the
|
42 |
+
denoising loop.
|
43 |
+
pred_original_sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` for images):
|
44 |
+
The predicted denoised sample `(x_{0})` based on the model output from the current timestep.
|
45 |
+
`pred_original_sample` can be used to preview progress or for guidance.
|
46 |
+
"""
|
47 |
+
|
48 |
+
prev_sample: torch.FloatTensor
|
49 |
+
denoised: Optional[torch.FloatTensor] = None
|
50 |
+
|
51 |
+
|
52 |
+
# Copied from diffusers.schedulers.scheduling_ddpm.betas_for_alpha_bar
|
53 |
+
def betas_for_alpha_bar(
|
54 |
+
num_diffusion_timesteps,
|
55 |
+
max_beta=0.999,
|
56 |
+
alpha_transform_type="cosine",
|
57 |
+
):
|
58 |
+
"""
|
59 |
+
Create a beta schedule that discretizes the given alpha_t_bar function, which defines the cumulative product of
|
60 |
+
(1-beta) over time from t = [0,1].
|
61 |
+
|
62 |
+
Contains a function alpha_bar that takes an argument t and transforms it to the cumulative product of (1-beta) up
|
63 |
+
to that part of the diffusion process.
|
64 |
+
|
65 |
+
|
66 |
+
Args:
|
67 |
+
num_diffusion_timesteps (`int`): the number of betas to produce.
|
68 |
+
max_beta (`float`): the maximum beta to use; use values lower than 1 to
|
69 |
+
prevent singularities.
|
70 |
+
alpha_transform_type (`str`, *optional*, default to `cosine`): the type of noise schedule for alpha_bar.
|
71 |
+
Choose from `cosine` or `exp`
|
72 |
+
|
73 |
+
Returns:
|
74 |
+
betas (`np.ndarray`): the betas used by the scheduler to step the model outputs
|
75 |
+
"""
|
76 |
+
if alpha_transform_type == "cosine":
|
77 |
+
|
78 |
+
def alpha_bar_fn(t):
|
79 |
+
return math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
|
80 |
+
|
81 |
+
elif alpha_transform_type == "exp":
|
82 |
+
|
83 |
+
def alpha_bar_fn(t):
|
84 |
+
return math.exp(t * -12.0)
|
85 |
+
|
86 |
+
else:
|
87 |
+
raise ValueError(f"Unsupported alpha_tranform_type: {alpha_transform_type}")
|
88 |
+
|
89 |
+
betas = []
|
90 |
+
for i in range(num_diffusion_timesteps):
|
91 |
+
t1 = i / num_diffusion_timesteps
|
92 |
+
t2 = (i + 1) / num_diffusion_timesteps
|
93 |
+
betas.append(min(1 - alpha_bar_fn(t2) / alpha_bar_fn(t1), max_beta))
|
94 |
+
return torch.tensor(betas, dtype=torch.float32)
|
95 |
+
|
96 |
+
|
97 |
+
# Copied from diffusers.schedulers.scheduling_ddim.rescale_zero_terminal_snr
|
98 |
+
def rescale_zero_terminal_snr(betas: torch.FloatTensor) -> torch.FloatTensor:
|
99 |
+
"""
|
100 |
+
Rescales betas to have zero terminal SNR Based on https://arxiv.org/pdf/2305.08891.pdf (Algorithm 1)
|
101 |
+
|
102 |
+
|
103 |
+
Args:
|
104 |
+
betas (`torch.FloatTensor`):
|
105 |
+
the betas that the scheduler is being initialized with.
|
106 |
+
|
107 |
+
Returns:
|
108 |
+
`torch.FloatTensor`: rescaled betas with zero terminal SNR
|
109 |
+
"""
|
110 |
+
# Convert betas to alphas_bar_sqrt
|
111 |
+
alphas = 1.0 - betas
|
112 |
+
alphas_cumprod = torch.cumprod(alphas, dim=0)
|
113 |
+
alphas_bar_sqrt = alphas_cumprod.sqrt()
|
114 |
+
|
115 |
+
# Store old values.
|
116 |
+
alphas_bar_sqrt_0 = alphas_bar_sqrt[0].clone()
|
117 |
+
alphas_bar_sqrt_T = alphas_bar_sqrt[-1].clone()
|
118 |
+
|
119 |
+
# Shift so the last timestep is zero.
|
120 |
+
alphas_bar_sqrt -= alphas_bar_sqrt_T
|
121 |
+
|
122 |
+
# Scale so the first timestep is back to the old value.
|
123 |
+
alphas_bar_sqrt *= alphas_bar_sqrt_0 / (alphas_bar_sqrt_0 - alphas_bar_sqrt_T)
|
124 |
+
|
125 |
+
# Convert alphas_bar_sqrt to betas
|
126 |
+
alphas_bar = alphas_bar_sqrt**2 # Revert sqrt
|
127 |
+
alphas = alphas_bar[1:] / alphas_bar[:-1] # Revert cumprod
|
128 |
+
alphas = torch.cat([alphas_bar[0:1], alphas])
|
129 |
+
betas = 1 - alphas
|
130 |
+
|
131 |
+
return betas
|
132 |
+
|
133 |
+
|
134 |
+
class LCMScheduler(SchedulerMixin, ConfigMixin):
|
135 |
+
"""
|
136 |
+
`LCMScheduler` extends the denoising procedure introduced in denoising diffusion probabilistic models (DDPMs) with
|
137 |
+
non-Markovian guidance.
|
138 |
+
|
139 |
+
This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. [`~ConfigMixin`] takes care of storing all config
|
140 |
+
attributes that are passed in the scheduler's `__init__` function, such as `num_train_timesteps`. They can be
|
141 |
+
accessed via `scheduler.config.num_train_timesteps`. [`SchedulerMixin`] provides general loading and saving
|
142 |
+
functionality via the [`SchedulerMixin.save_pretrained`] and [`~SchedulerMixin.from_pretrained`] functions.
|
143 |
+
|
144 |
+
Args:
|
145 |
+
num_train_timesteps (`int`, defaults to 1000):
|
146 |
+
The number of diffusion steps to train the model.
|
147 |
+
beta_start (`float`, defaults to 0.0001):
|
148 |
+
The starting `beta` value of inference.
|
149 |
+
beta_end (`float`, defaults to 0.02):
|
150 |
+
The final `beta` value.
|
151 |
+
beta_schedule (`str`, defaults to `"linear"`):
|
152 |
+
The beta schedule, a mapping from a beta range to a sequence of betas for stepping the model. Choose from
|
153 |
+
`linear`, `scaled_linear`, or `squaredcos_cap_v2`.
|
154 |
+
trained_betas (`np.ndarray`, *optional*):
|
155 |
+
Pass an array of betas directly to the constructor to bypass `beta_start` and `beta_end`.
|
156 |
+
original_inference_steps (`int`, *optional*, defaults to 50):
|
157 |
+
The default number of inference steps used to generate a linearly-spaced timestep schedule, from which we
|
158 |
+
will ultimately take `num_inference_steps` evenly spaced timesteps to form the final timestep schedule.
|
159 |
+
clip_sample (`bool`, defaults to `True`):
|
160 |
+
Clip the predicted sample for numerical stability.
|
161 |
+
clip_sample_range (`float`, defaults to 1.0):
|
162 |
+
The maximum magnitude for sample clipping. Valid only when `clip_sample=True`.
|
163 |
+
set_alpha_to_one (`bool`, defaults to `True`):
|
164 |
+
Each diffusion step uses the alphas product value at that step and at the previous one. For the final step
|
165 |
+
there is no previous alpha. When this option is `True` the previous alpha product is fixed to `1`,
|
166 |
+
otherwise it uses the alpha value at step 0.
|
167 |
+
steps_offset (`int`, defaults to 0):
|
168 |
+
An offset added to the inference steps. You can use a combination of `offset=1` and
|
169 |
+
`set_alpha_to_one=False` to make the last step use step 0 for the previous alpha product like in Stable
|
170 |
+
Diffusion.
|
171 |
+
prediction_type (`str`, defaults to `epsilon`, *optional*):
|
172 |
+
Prediction type of the scheduler function; can be `epsilon` (predicts the noise of the diffusion process),
|
173 |
+
`sample` (directly predicts the noisy sample`) or `v_prediction` (see section 2.4 of [Imagen
|
174 |
+
Video](https://imagen.research.google/video/paper.pdf) paper).
|
175 |
+
thresholding (`bool`, defaults to `False`):
|
176 |
+
Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
|
177 |
+
as Stable Diffusion.
|
178 |
+
dynamic_thresholding_ratio (`float`, defaults to 0.995):
|
179 |
+
The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
|
180 |
+
sample_max_value (`float`, defaults to 1.0):
|
181 |
+
The threshold value for dynamic thresholding. Valid only when `thresholding=True`.
|
182 |
+
timestep_spacing (`str`, defaults to `"leading"`):
|
183 |
+
The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
|
184 |
+
Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
|
185 |
+
rescale_betas_zero_snr (`bool`, defaults to `False`):
|
186 |
+
Whether to rescale the betas to have zero terminal SNR. This enables the model to generate very bright and
|
187 |
+
dark samples instead of limiting it to samples with medium brightness. Loosely related to
|
188 |
+
[`--offset_noise`](https://github.com/huggingface/diffusers/blob/74fd735eb073eb1d774b1ab4154a0876eb82f055/examples/dreambooth/train_dreambooth.py#L506).
|
189 |
+
"""
|
190 |
+
|
191 |
+
order = 1
|
192 |
+
|
193 |
+
@register_to_config
|
194 |
+
def __init__(
|
195 |
+
self,
|
196 |
+
num_train_timesteps: int = 1000,
|
197 |
+
beta_start: float = 0.00085,
|
198 |
+
beta_end: float = 0.012,
|
199 |
+
beta_schedule: str = "scaled_linear",
|
200 |
+
trained_betas: Optional[Union[np.ndarray, List[float]]] = None,
|
201 |
+
original_inference_steps: int = 50,
|
202 |
+
clip_sample: bool = False,
|
203 |
+
clip_sample_range: float = 1.0,
|
204 |
+
set_alpha_to_one: bool = True,
|
205 |
+
steps_offset: int = 0,
|
206 |
+
prediction_type: str = "epsilon",
|
207 |
+
thresholding: bool = False,
|
208 |
+
dynamic_thresholding_ratio: float = 0.995,
|
209 |
+
sample_max_value: float = 1.0,
|
210 |
+
timestep_spacing: str = "leading",
|
211 |
+
rescale_betas_zero_snr: bool = False,
|
212 |
+
):
|
213 |
+
if trained_betas is not None:
|
214 |
+
self.betas = torch.tensor(trained_betas, dtype=torch.float32)
|
215 |
+
elif beta_schedule == "linear":
|
216 |
+
self.betas = torch.linspace(beta_start, beta_end, num_train_timesteps, dtype=torch.float32)
|
217 |
+
elif beta_schedule == "scaled_linear":
|
218 |
+
# this schedule is very specific to the latent diffusion model.
|
219 |
+
self.betas = (
|
220 |
+
torch.linspace(beta_start**0.5, beta_end**0.5, num_train_timesteps, dtype=torch.float32) ** 2
|
221 |
+
)
|
222 |
+
elif beta_schedule == "squaredcos_cap_v2":
|
223 |
+
# Glide cosine schedule
|
224 |
+
self.betas = betas_for_alpha_bar(num_train_timesteps)
|
225 |
+
else:
|
226 |
+
raise NotImplementedError(f"{beta_schedule} does is not implemented for {self.__class__}")
|
227 |
+
|
228 |
+
# Rescale for zero SNR
|
229 |
+
if rescale_betas_zero_snr:
|
230 |
+
self.betas = rescale_zero_terminal_snr(self.betas)
|
231 |
+
|
232 |
+
self.alphas = 1.0 - self.betas
|
233 |
+
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
|
234 |
+
|
235 |
+
# At every step in ddim, we are looking into the previous alphas_cumprod
|
236 |
+
# For the final step, there is no previous alphas_cumprod because we are already at 0
|
237 |
+
# `set_alpha_to_one` decides whether we set this parameter simply to one or
|
238 |
+
# whether we use the final alpha of the "non-previous" one.
|
239 |
+
self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]
|
240 |
+
|
241 |
+
# standard deviation of the initial noise distribution
|
242 |
+
self.init_noise_sigma = 1.0
|
243 |
+
|
244 |
+
# setable values
|
245 |
+
self.num_inference_steps = None
|
246 |
+
self.timesteps = torch.from_numpy(np.arange(0, num_train_timesteps)[::-1].copy().astype(np.int64))
|
247 |
+
|
248 |
+
self._step_index = None
|
249 |
+
|
250 |
+
# Copied from diffusers.schedulers.scheduling_euler_discrete.EulerDiscreteScheduler._init_step_index
|
251 |
+
def _init_step_index(self, timestep):
|
252 |
+
if isinstance(timestep, torch.Tensor):
|
253 |
+
timestep = timestep.to(self.timesteps.device)
|
254 |
+
|
255 |
+
index_candidates = (self.timesteps == timestep).nonzero()
|
256 |
+
|
257 |
+
# The sigma index that is taken for the **very** first `step`
|
258 |
+
# is always the second index (or the last index if there is only 1)
|
259 |
+
# This way we can ensure we don't accidentally skip a sigma in
|
260 |
+
# case we start in the middle of the denoising schedule (e.g. for image-to-image)
|
261 |
+
if len(index_candidates) > 1:
|
262 |
+
step_index = index_candidates[1]
|
263 |
+
else:
|
264 |
+
step_index = index_candidates[0]
|
265 |
+
|
266 |
+
self._step_index = step_index.item()
|
267 |
+
|
268 |
+
@property
|
269 |
+
def step_index(self):
|
270 |
+
return self._step_index
|
271 |
+
|
272 |
+
def scale_model_input(self, sample: torch.FloatTensor, timestep: Optional[int] = None) -> torch.FloatTensor:
|
273 |
+
"""
|
274 |
+
Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
|
275 |
+
current timestep.
|
276 |
+
|
277 |
+
Args:
|
278 |
+
sample (`torch.FloatTensor`):
|
279 |
+
The input sample.
|
280 |
+
timestep (`int`, *optional*):
|
281 |
+
The current timestep in the diffusion chain.
|
282 |
+
Returns:
|
283 |
+
`torch.FloatTensor`:
|
284 |
+
A scaled input sample.
|
285 |
+
"""
|
286 |
+
return sample
|
287 |
+
|
288 |
+
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
|
289 |
+
def _threshold_sample(self, sample: torch.FloatTensor) -> torch.FloatTensor:
|
290 |
+
"""
|
291 |
+
"Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
|
292 |
+
prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
|
293 |
+
s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
|
294 |
+
pixels from saturation at each step. We find that dynamic thresholding results in significantly better
|
295 |
+
photorealism as well as better image-text alignment, especially when using very large guidance weights."
|
296 |
+
|
297 |
+
https://arxiv.org/abs/2205.11487
|
298 |
+
"""
|
299 |
+
dtype = sample.dtype
|
300 |
+
batch_size, channels, *remaining_dims = sample.shape
|
301 |
+
|
302 |
+
if dtype not in (torch.float32, torch.float64):
|
303 |
+
sample = sample.float() # upcast for quantile calculation, and clamp not implemented for cpu half
|
304 |
+
|
305 |
+
# Flatten sample for doing quantile calculation along each image
|
306 |
+
sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
|
307 |
+
|
308 |
+
abs_sample = sample.abs() # "a certain percentile absolute pixel value"
|
309 |
+
|
310 |
+
s = torch.quantile(abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
|
311 |
+
s = torch.clamp(
|
312 |
+
s, min=1, max=self.config.sample_max_value
|
313 |
+
) # When clamped to min=1, equivalent to standard clipping to [-1, 1]
|
314 |
+
s = s.unsqueeze(1) # (batch_size, 1) because clamp will broadcast along dim=0
|
315 |
+
sample = torch.clamp(sample, -s, s) / s # "we threshold xt0 to the range [-s, s] and then divide by s"
|
316 |
+
|
317 |
+
sample = sample.reshape(batch_size, channels, *remaining_dims)
|
318 |
+
sample = sample.to(dtype)
|
319 |
+
|
320 |
+
return sample
|
321 |
+
|
322 |
+
def set_timesteps(
|
323 |
+
self,
|
324 |
+
num_inference_steps: int,
|
325 |
+
device: Union[str, torch.device] = None,
|
326 |
+
original_inference_steps: Optional[int] = None,
|
327 |
+
):
|
328 |
+
"""
|
329 |
+
Sets the discrete timesteps used for the diffusion chain (to be run before inference).
|
330 |
+
|
331 |
+
Args:
|
332 |
+
num_inference_steps (`int`):
|
333 |
+
The number of diffusion steps used when generating samples with a pre-trained model.
|
334 |
+
device (`str` or `torch.device`, *optional*):
|
335 |
+
The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
|
336 |
+
original_inference_steps (`int`, *optional*):
|
337 |
+
The original number of inference steps, which will be used to generate a linearly-spaced timestep
|
338 |
+
schedule (which is different from the standard `diffusers` implementation). We will then take
|
339 |
+
`num_inference_steps` timesteps from this schedule, evenly spaced in terms of indices, and use that as
|
340 |
+
our final timestep schedule. If not set, this will default to the `original_inference_steps` attribute.
|
341 |
+
"""
|
342 |
+
|
343 |
+
if num_inference_steps > self.config.num_train_timesteps:
|
344 |
+
raise ValueError(
|
345 |
+
f"`num_inference_steps`: {num_inference_steps} cannot be larger than `self.config.train_timesteps`:"
|
346 |
+
f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
|
347 |
+
f" maximal {self.config.num_train_timesteps} timesteps."
|
348 |
+
)
|
349 |
+
|
350 |
+
self.num_inference_steps = num_inference_steps
|
351 |
+
original_steps = (
|
352 |
+
original_inference_steps if original_inference_steps is not None else self.original_inference_steps
|
353 |
+
)
|
354 |
+
|
355 |
+
if original_steps > self.config.num_train_timesteps:
|
356 |
+
raise ValueError(
|
357 |
+
f"`original_steps`: {original_steps} cannot be larger than `self.config.train_timesteps`:"
|
358 |
+
f" {self.config.num_train_timesteps} as the unet model trained with this scheduler can only handle"
|
359 |
+
f" maximal {self.config.num_train_timesteps} timesteps."
|
360 |
+
)
|
361 |
+
|
362 |
+
if num_inference_steps > original_steps:
|
363 |
+
raise ValueError(
|
364 |
+
f"`num_inference_steps`: {num_inference_steps} cannot be larger than `original_inference_steps`:"
|
365 |
+
f" {original_steps} because the final timestep schedule will be a subset of the"
|
366 |
+
f" `original_inference_steps`-sized initial timestep schedule."
|
367 |
+
)
|
368 |
+
|
369 |
+
# LCM Timesteps Setting
|
370 |
+
# Currently, only linear spacing is supported.
|
371 |
+
c = self.config.num_train_timesteps // original_steps
|
372 |
+
# LCM Training Steps Schedule
|
373 |
+
lcm_origin_timesteps = np.asarray(list(range(1, original_steps + 1))) * c - 1
|
374 |
+
skipping_step = len(lcm_origin_timesteps) // num_inference_steps
|
375 |
+
# LCM Inference Steps Schedule
|
376 |
+
timesteps = lcm_origin_timesteps[::-skipping_step][:num_inference_steps]
|
377 |
+
|
378 |
+
self.timesteps = torch.from_numpy(timesteps.copy()).to(device=device, dtype=torch.long)
|
379 |
+
|
380 |
+
self._step_index = None
|
381 |
+
|
382 |
+
def get_scalings_for_boundary_condition_discrete(self, t):
|
383 |
+
self.sigma_data = 0.5 # Default: 0.5
|
384 |
+
|
385 |
+
# By dividing 0.1: This is almost a delta function at t=0.
|
386 |
+
c_skip = self.sigma_data**2 / ((t / 0.1) ** 2 + self.sigma_data**2)
|
387 |
+
c_out = (t / 0.1) / ((t / 0.1) ** 2 + self.sigma_data**2) ** 0.5
|
388 |
+
return c_skip, c_out
|
389 |
+
|
390 |
+
def step(
|
391 |
+
self,
|
392 |
+
model_output: torch.FloatTensor,
|
393 |
+
timestep: int,
|
394 |
+
sample: torch.FloatTensor,
|
395 |
+
generator: Optional[torch.Generator] = None,
|
396 |
+
return_dict: bool = True,
|
397 |
+
) -> Union[LCMSchedulerOutput, Tuple]:
|
398 |
+
"""
|
399 |
+
Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion
|
400 |
+
process from the learned model outputs (most often the predicted noise).
|
401 |
+
|
402 |
+
Args:
|
403 |
+
model_output (`torch.FloatTensor`):
|
404 |
+
The direct output from learned diffusion model.
|
405 |
+
timestep (`float`):
|
406 |
+
The current discrete timestep in the diffusion chain.
|
407 |
+
sample (`torch.FloatTensor`):
|
408 |
+
A current instance of a sample created by the diffusion process.
|
409 |
+
generator (`torch.Generator`, *optional*):
|
410 |
+
A random number generator.
|
411 |
+
return_dict (`bool`, *optional*, defaults to `True`):
|
412 |
+
Whether or not to return a [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] or `tuple`.
|
413 |
+
Returns:
|
414 |
+
[`~schedulers.scheduling_utils.LCMSchedulerOutput`] or `tuple`:
|
415 |
+
If return_dict is `True`, [`~schedulers.scheduling_lcm.LCMSchedulerOutput`] is returned, otherwise a
|
416 |
+
tuple is returned where the first element is the sample tensor.
|
417 |
+
"""
|
418 |
+
if self.num_inference_steps is None:
|
419 |
+
raise ValueError(
|
420 |
+
"Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
|
421 |
+
)
|
422 |
+
|
423 |
+
if self.step_index is None:
|
424 |
+
self._init_step_index(timestep)
|
425 |
+
|
426 |
+
# 1. get previous step value
|
427 |
+
prev_step_index = self.step_index + 1
|
428 |
+
if prev_step_index < len(self.timesteps):
|
429 |
+
prev_timestep = self.timesteps[prev_step_index]
|
430 |
+
else:
|
431 |
+
prev_timestep = timestep
|
432 |
+
|
433 |
+
# 2. compute alphas, betas
|
434 |
+
alpha_prod_t = self.alphas_cumprod[timestep]
|
435 |
+
alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod
|
436 |
+
|
437 |
+
beta_prod_t = 1 - alpha_prod_t
|
438 |
+
beta_prod_t_prev = 1 - alpha_prod_t_prev
|
439 |
+
|
440 |
+
# 3. Get scalings for boundary conditions
|
441 |
+
c_skip, c_out = self.get_scalings_for_boundary_condition_discrete(timestep)
|
442 |
+
|
443 |
+
# 4. Compute the predicted original sample x_0 based on the model parameterization
|
444 |
+
if self.config.prediction_type == "epsilon": # noise-prediction
|
445 |
+
predicted_original_sample = (sample - beta_prod_t.sqrt() * model_output) / alpha_prod_t.sqrt()
|
446 |
+
elif self.config.prediction_type == "sample": # x-prediction
|
447 |
+
predicted_original_sample = model_output
|
448 |
+
elif self.config.prediction_type == "v_prediction": # v-prediction
|
449 |
+
predicted_original_sample = alpha_prod_t.sqrt() * sample - beta_prod_t.sqrt() * model_output
|
450 |
+
else:
|
451 |
+
raise ValueError(
|
452 |
+
f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` or"
|
453 |
+
" `v_prediction` for `LCMScheduler`."
|
454 |
+
)
|
455 |
+
|
456 |
+
# 5. Clip or threshold "predicted x_0"
|
457 |
+
if self.config.thresholding:
|
458 |
+
predicted_original_sample = self._threshold_sample(predicted_original_sample)
|
459 |
+
elif self.config.clip_sample:
|
460 |
+
predicted_original_sample = predicted_original_sample.clamp(
|
461 |
+
-self.config.clip_sample_range, self.config.clip_sample_range
|
462 |
+
)
|
463 |
+
|
464 |
+
# 6. Denoise model output using boundary conditions
|
465 |
+
denoised = c_out * predicted_original_sample + c_skip * sample
|
466 |
+
|
467 |
+
# 7. Sample and inject noise z ~ N(0, I) for MultiStep Inference
|
468 |
+
# Noise is not used for one-step sampling.
|
469 |
+
if len(self.timesteps) > 1:
|
470 |
+
noise = randn_tensor(model_output.shape, generator=generator, device=model_output.device)
|
471 |
+
prev_sample = alpha_prod_t_prev.sqrt() * denoised + beta_prod_t_prev.sqrt() * noise
|
472 |
+
else:
|
473 |
+
prev_sample = denoised
|
474 |
+
|
475 |
+
# upon completion increase step index by one
|
476 |
+
self._step_index += 1
|
477 |
+
|
478 |
+
if not return_dict:
|
479 |
+
return (prev_sample, denoised)
|
480 |
+
|
481 |
+
return LCMSchedulerOutput(prev_sample=prev_sample, denoised=denoised)
|
482 |
+
|
483 |
+
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.add_noise
|
484 |
+
def add_noise(
|
485 |
+
self,
|
486 |
+
original_samples: torch.FloatTensor,
|
487 |
+
noise: torch.FloatTensor,
|
488 |
+
timesteps: torch.IntTensor,
|
489 |
+
) -> torch.FloatTensor:
|
490 |
+
# Make sure alphas_cumprod and timestep have same device and dtype as original_samples
|
491 |
+
alphas_cumprod = self.alphas_cumprod.to(device=original_samples.device, dtype=original_samples.dtype)
|
492 |
+
timesteps = timesteps.to(original_samples.device)
|
493 |
+
|
494 |
+
sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
|
495 |
+
sqrt_alpha_prod = sqrt_alpha_prod.flatten()
|
496 |
+
while len(sqrt_alpha_prod.shape) < len(original_samples.shape):
|
497 |
+
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
|
498 |
+
|
499 |
+
sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
|
500 |
+
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
|
501 |
+
while len(sqrt_one_minus_alpha_prod.shape) < len(original_samples.shape):
|
502 |
+
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
|
503 |
+
|
504 |
+
noisy_samples = sqrt_alpha_prod * original_samples + sqrt_one_minus_alpha_prod * noise
|
505 |
+
return noisy_samples
|
506 |
+
|
507 |
+
# Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler.get_velocity
|
508 |
+
def get_velocity(
|
509 |
+
self, sample: torch.FloatTensor, noise: torch.FloatTensor, timesteps: torch.IntTensor
|
510 |
+
) -> torch.FloatTensor:
|
511 |
+
# Make sure alphas_cumprod and timestep have same device and dtype as sample
|
512 |
+
alphas_cumprod = self.alphas_cumprod.to(device=sample.device, dtype=sample.dtype)
|
513 |
+
timesteps = timesteps.to(sample.device)
|
514 |
+
|
515 |
+
sqrt_alpha_prod = alphas_cumprod[timesteps] ** 0.5
|
516 |
+
sqrt_alpha_prod = sqrt_alpha_prod.flatten()
|
517 |
+
while len(sqrt_alpha_prod.shape) < len(sample.shape):
|
518 |
+
sqrt_alpha_prod = sqrt_alpha_prod.unsqueeze(-1)
|
519 |
+
|
520 |
+
sqrt_one_minus_alpha_prod = (1 - alphas_cumprod[timesteps]) ** 0.5
|
521 |
+
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.flatten()
|
522 |
+
while len(sqrt_one_minus_alpha_prod.shape) < len(sample.shape):
|
523 |
+
sqrt_one_minus_alpha_prod = sqrt_one_minus_alpha_prod.unsqueeze(-1)
|
524 |
+
|
525 |
+
velocity = sqrt_alpha_prod * noise - sqrt_one_minus_alpha_prod * sample
|
526 |
+
return velocity
|
527 |
+
|
528 |
+
def __len__(self):
|
529 |
+
return self.config.num_train_timesteps
|
requirements.txt
ADDED
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
accelerate
|
2 |
+
diffusers
|
3 |
+
gradio==3.48.0
|
4 |
+
Pillow
|
5 |
+
--extra-index-url https://download.pytorch.org/whl/cpu
|
6 |
+
torch==2.0.1
|
7 |
+
openvino==2023.1.0
|
8 |
+
optimum-intel==1.11.0
|
9 |
+
transformers
|
10 |
+
opencv-python
|
11 |
+
|
12 |
+
git+https://huggingface.co/spaces/Wauplin/gradio-user-history
|
style.css
ADDED
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
h1 {
|
2 |
+
text-align: center;
|
3 |
+
}
|
4 |
+
|
5 |
+
#duplicate-button {
|
6 |
+
margin: auto;
|
7 |
+
color: #fff;
|
8 |
+
background: #1565c0;
|
9 |
+
border-radius: 100vh;
|
10 |
+
}
|
11 |
+
|
12 |
+
#component-0 {
|
13 |
+
max-width: 830px;
|
14 |
+
margin: auto;
|
15 |
+
padding-top: 1.5rem;
|
16 |
+
}
|