deinferno commited on
Commit
5b1f3e5
1 Parent(s): 515cc24

Upload space

Browse files
Files changed (7) hide show
  1. LICENSE +21 -0
  2. README.md +9 -6
  3. app.py +176 -0
  4. lcm_ov_pipeline.py +388 -0
  5. lcm_scheduler.py +529 -0
  6. requirements.txt +12 -0
  7. 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 Model OpenVino CPU
3
- emoji: 👀
4
- colorFrom: gray
5
- colorTo: blue
6
  sdk: gradio
7
- sdk_version: 3.50.2
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
+ }