nodes/aszc-dev-ComfyUI-CoreMLSuite
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Remove Simple LCM Sampler

Browse Source
This commit is contained in:
aszc-dev
2023-11-03 01:29:50 +01:00
parent c01c60e3c1
commit 1ebd9e72ae
5 changed files with 4 additions and 382 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+2 -5
View File
@@ -7,7 +7,6 @@ from coreml_suite.nodes import CoreMLLoaderUNet, CoreMLSampler, CoreMLModelAdapt
from coreml_suite.lcm import (
CoreMLSamplerLCM,
CoreMLConverterLCM,
CoreMLSamplerLCM_Simple,
)
NODE_CLASS_MAPPINGS = {
@@ -15,14 +14,12 @@ NODE_CLASS_MAPPINGS = {
"CoreMLSampler": CoreMLSampler,
"CoreMLModelAdapter": CoreMLModelAdapter,
"Core ML LCM Sampler": CoreMLSamplerLCM,
"Core ML LCM Sampler (Simple)": CoreMLSamplerLCM_Simple,
"CoreMLConverterLCM": CoreMLConverterLCM,
"Core ML LCM Converter": CoreMLConverterLCM,
}
NODE_DISPLAY_NAME_MAPPINGS = {
"CoreMLUNetLoader": "Load Core ML UNet",
"CoreMLSampler": "Core ML Sampler",
"CoreMLModelAdapter": "Core ML Adapter (Experimental)",
"Core ML LCM Sampler": "Core ML LCM Sampler",
"Core ML LCM Sampler (Simple)": "Core ML LCM Sampler (Simple)",
"CoreMLConverterLCM": "Convert LCM to Core ML",
"Core ML LCM Converter": "Convert LCM to Core ML",
}
coreml_suite/lcm/__init__.py
+2 -2
View File
@@ -1,4 +1,4 @@
from .lcm_sampler import CoreMLSamplerLCM, CoreMLSamplerLCM_Simple
from .lcm_sampler import CoreMLSamplerLCM
from .nodes import CoreMLConverterLCM
__all__ = ["CoreMLSamplerLCM", "CoreMLSamplerLCM_Simple", "CoreMLConverterLCM"]
__all__ = ["CoreMLSamplerLCM", "CoreMLConverterLCM"]
coreml_suite/lcm/lcm_pipeline.py
-292
View File
@@ -1,292 +0,0 @@
import torch
from diffusers import DiffusionPipeline, AutoencoderKL, UNet2DConditionModel
from transformers import CLIPTokenizer, CLIPTextModel, CLIPImageProcessor
from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
from diffusers.image_processor import VaeImageProcessor
from typing import List, Optional, Union, Dict, Any
from comfy.model_management import get_torch_device
# from diffusers import logging
# logger = logging.get_logger(__name__) # pylint: disable=invalid-name
class LatentConsistencyModelPipeline(DiffusionPipeline):
def __init__(
self,
vae: AutoencoderKL,
text_encoder: CLIPTextModel,
tokenizer: CLIPTokenizer,
unet: UNet2DConditionModel,
scheduler: None,
safety_checker: None,
feature_extractor: CLIPImageProcessor,
):
super().__init__()
self.register_modules(
vae=vae,
text_encoder=text_encoder,
tokenizer=tokenizer,
unet=unet,
scheduler=scheduler,
safety_checker=safety_checker,
feature_extractor=feature_extractor,
)
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
def _encode_prompt(
self,
prompt,
device,
num_images_per_prompt,
prompt_embeds: None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
prompt (`str` or `List[str]`, *optional*):
prompt to be encoded
device: (`torch.device`):
torch device
num_images_per_prompt (`int`):
number of images that should be generated per prompt
prompt_embeds (`torch.FloatTensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
"""
if prompt is not None and isinstance(prompt, str):
batch_size = 1
elif prompt is not None and isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
if prompt_embeds is None:
text_inputs = self.tokenizer(
prompt,
padding="max_length",
max_length=self.tokenizer.model_max_length,
truncation=True,
return_tensors="pt",
)
text_input_ids = text_inputs.input_ids
untruncated_ids = self.tokenizer(
prompt, padding="longest", return_tensors="pt"
).input_ids
if untruncated_ids.shape[-1] >= text_input_ids.shape[
-1
] and not torch.equal(text_input_ids, untruncated_ids):
removed_text = self.tokenizer.batch_decode(
untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
)
print(
"The following part of your input was truncated because CLIP can only handle sequences up to"
f" {self.tokenizer.model_max_length} tokens: {removed_text}"
)
if (
hasattr(self.text_encoder.config, "use_attention_mask")
and self.text_encoder.config.use_attention_mask
):
attention_mask = text_inputs.attention_mask.to(device)
else:
attention_mask = None
prompt_embeds = self.text_encoder(
text_input_ids.to(device),
attention_mask=attention_mask,
)
prompt_embeds = prompt_embeds[0]
if self.text_encoder is not None:
prompt_embeds_dtype = self.text_encoder.dtype
elif self.unet is not None:
prompt_embeds_dtype = self.unet.dtype
else:
prompt_embeds_dtype = prompt_embeds.dtype
prompt_embeds = prompt_embeds.to(dtype=prompt_embeds_dtype, device=device)
bs_embed, seq_len, _ = prompt_embeds.shape
# duplicate text embeddings for each generation per prompt, using mps friendly method
prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
prompt_embeds = prompt_embeds.view(
bs_embed * num_images_per_prompt, seq_len, -1
)
# Don't need to get uncond prompt embedding because of LCM Guided Distillation
return prompt_embeds
# ¯\_(ツ)_/¯
def run_safety_checker(self, image, device, dtype):
return image, None
def prepare_latents(
self,
batch_size,
num_channels_latents,
height,
width,
dtype,
device,
latents=None,
):
shape = (
batch_size,
num_channels_latents,
height // self.vae_scale_factor,
width // self.vae_scale_factor,
)
if latents is None:
latents = torch.randn(shape, dtype=dtype).to(device)
else:
latents = latents.to(device)
# scale the initial noise by the standard deviation required by the scheduler
latents = latents * self.scheduler.init_noise_sigma
return latents
def get_w_embedding(self, w, embedding_dim=512, dtype=torch.float32):
"""
see https://github.com/google-research/vdm/blob/dc27b98a554f65cdc654b800da5aa1846545d41b/model_vdm.py#L298
Args:
timesteps: torch.Tensor: generate embedding vectors at these timesteps
embedding_dim: int: dimension of the embeddings to generate
dtype: data type of the generated embeddings
Returns:
embedding vectors with shape `(len(timesteps), embedding_dim)`
"""
assert len(w.shape) == 1
w = w * 1000.0
half_dim = embedding_dim // 2
emb = torch.log(torch.tensor(10000.0)) / (half_dim - 1)
emb = torch.exp(torch.arange(half_dim, dtype=dtype) * -emb)
emb = w.to(dtype)[:, None] * emb[None, :]
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
if embedding_dim % 2 == 1: # zero pad
emb = torch.nn.functional.pad(emb, (0, 1))
assert emb.shape == (w.shape[0], embedding_dim)
return emb
@torch.no_grad()
def __call__(
self,
prompt: Union[str, List[str]] = None,
height: Optional[int] = 768,
width: Optional[int] = 768,
guidance_scale: float = 7.5,
num_images_per_prompt: Optional[int] = 1,
latents: Optional[torch.FloatTensor] = None,
num_inference_steps: int = 4,
lcm_origin_steps: int = 50,
prompt_embeds: Optional[torch.FloatTensor] = None,
output_type: Optional[str] = "pil",
return_dict: bool = True,
cross_attention_kwargs: Optional[Dict[str, Any]] = None,
):
# 0. Default height and width to unet
height = height or self.unet.config.sample_size * self.vae_scale_factor
width = width or self.unet.config.sample_size * self.vae_scale_factor
# 2. Define call parameters
if prompt is not None and isinstance(prompt, str):
batch_size = 1
elif prompt is not None and isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
device = get_torch_device()
# do_classifier_free_guidance = guidance_scale > 0.0 # In LCM Implementation: cfg_noise = noise_cond + cfg_scale * (noise_cond - noise_uncond) , (cfg_scale > 0.0 using CFG)
# 3. Encode input prompt
prompt_embeds = self._encode_prompt(
prompt,
device,
num_images_per_prompt,
prompt_embeds=prompt_embeds,
)
# 4. Prepare timesteps
self.scheduler.set_timesteps(num_inference_steps, lcm_origin_steps)
timesteps = self.scheduler.timesteps
# 5. Prepare latent variable
num_channels_latents = self.unet.config.in_channels
latents = self.prepare_latents(
batch_size * num_images_per_prompt,
num_channels_latents,
height,
width,
prompt_embeds.dtype,
device,
latents,
)
bs = batch_size * num_images_per_prompt
# 6. Get Guidance Scale Embedding
w = torch.tensor(guidance_scale).repeat(bs)
w_embedding = self.get_w_embedding(w, embedding_dim=256).to(
device=device, dtype=latents.dtype
)
# 7. LCM MultiStep Sampling Loop:
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
ts = torch.full((bs,), t, device=device, dtype=torch.long)
latents = latents.to(prompt_embeds.dtype)
# model prediction (v-prediction, eps, x)
print("latents", latents.shape)
model_pred = self.unet(
latents,
ts,
timestep_cond=w_embedding,
encoder_hidden_states=prompt_embeds,
cross_attention_kwargs=cross_attention_kwargs,
return_dict=False,
)[0]
print("model_pred", model_pred.shape)
# compute the previous noisy sample x_t -> x_t-1
latents, denoised = self.scheduler.step(
model_pred, i, t, latents, return_dict=False
)
# # call the callback, if provided
# if i == len(timesteps) - 1:
progress_bar.update()
denoised = denoised.to(prompt_embeds.dtype)
if not output_type == "latent":
image = self.vae.decode(
denoised / self.vae.config.scaling_factor, return_dict=False
)[0]
image, has_nsfw_concept = self.run_safety_checker(
image, device, prompt_embeds.dtype
)
else:
image = denoised
has_nsfw_concept = None
if has_nsfw_concept is None:
do_denormalize = [True] * image.shape[0]
else:
do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
image = self.image_processor.postprocess(
image, output_type=output_type, do_denormalize=do_denormalize
)
if not return_dict:
return (image, has_nsfw_concept)
return StableDiffusionPipelineOutput(
images=image, nsfw_content_detected=has_nsfw_concept
)
coreml_suite/lcm/lcm_sampler.py
-82
View File
@@ -5,97 +5,15 @@ import torch
from diffusers.utils.torch_utils import randn_tensor
from tqdm import tqdm
import comfy.utils
import latent_preview
from comfy.model_management import get_torch_device
from comfy.model_patcher import ModelPatcher
from coreml_suite.lcm.lcm_pipeline import LatentConsistencyModelPipeline
from coreml_suite.lcm.lcm_scheduler import LCMScheduler
from coreml_suite.logger import logger
from coreml_suite.models import get_model_config, CoreMLModelWrapperLCM
from coreml_suite.nodes import CoreMLSampler
class CoreMLSamplerLCM_Simple:
def __init__(self):
self.scheduler = LCMScheduler.from_pretrained(
os.path.join(os.path.dirname(__file__), "scheduler_config.json")
)
self.pipe = None
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"coreml_model": ("COREML_UNET",),
"seed": ("INT", {"default": 0, "min": 0, "max": 0xFFFFFFFFFFFFFFFF}),
"steps": ("INT", {"default": 4, "min": 1, "max": 10000}),
"cfg": (
"FLOAT",
{
"default": 8.0,
"min": 0.0,
"max": 100.0,
"step": 0.5,
"round": 0.01,
},
),
"num_images": ("INT", {"default": 1, "min": 1, "max": 64}),
"positive_prompt": ("STRING", {"multiline": True}),
}
}
RETURN_TYPES = ("IMAGE",)
FUNCTION = "sample"
CATEGORY = "sampling"
def sample(
self,
coreml_model,
seed,
steps,
cfg,
positive_prompt,
num_images,
):
height = coreml_model.expected_inputs["sample"]["shape"][2] * 8
width = coreml_model.expected_inputs["sample"]["shape"][3] * 8
model_config = get_model_config()
wrapped_model = CoreMLModelWrapperLCM(model_config, coreml_model)
if self.pipe is None:
self.pipe = LatentConsistencyModelPipeline.from_pretrained(
pretrained_model_name_or_path="SimianLuo/LCM_Dreamshaper_v7",
scheduler=self.scheduler,
safety_checker=None,
)
self.pipe.to(torch_device=get_torch_device(), torch_dtype=torch.float16)
coreml_unet = wrapped_model
coreml_unet.config = self.pipe.unet.config
self.pipe.unet = coreml_unet
torch.manual_seed(seed)
result = self.pipe(
prompt=positive_prompt,
width=width,
height=height,
guidance_scale=cfg,
num_inference_steps=steps,
num_images_per_prompt=num_images,
lcm_origin_steps=50,
output_type="np",
).images
images_tensor = torch.from_numpy(result)
return (images_tensor,)
class CoreMLSamplerLCM(CoreMLSampler):
@classmethod
def INPUT_TYPES(s):
coreml_suite/lcm/unet.py
-1
View File
@@ -1,4 +1,3 @@
from diffusers.configuration_utils import register_to_config
from overrides import overrides
from python_coreml_stable_diffusion.unet import UNet2DConditionModel, TimestepEmbedding