diff --git a/requirements_base.txt b/requirements_base.txt index 82b06e8..dcfba99 100644 --- a/requirements_base.txt +++ b/requirements_base.txt @@ -17,7 +17,6 @@ albumentations==1.4.15 albucore==0.0.16 pydantic omegaconf -k-diffusion open_clip_torch timm==1.0.22 prodigyopt diff --git a/toolkit/pipelines.py b/toolkit/pipelines.py index 3ccf6b3..3b92cab 100644 --- a/toolkit/pipelines.py +++ b/toolkit/pipelines.py @@ -1,5 +1,3 @@ -import importlib -import inspect from typing import Union, List, Optional, Dict, Any, Tuple, Callable import numpy as np @@ -8,12 +6,9 @@ from diffusers import StableDiffusionXLPipeline, StableDiffusionPipeline, LMSDis from diffusers.pipelines.flux.pipeline_flux import calculate_shift, retrieve_timesteps from diffusers.pipelines.flux.pipeline_output import FluxPipelineOutput from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput -# from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_k_diffusion import ModelWrapper from diffusers.pipelines.stable_diffusion_xl.pipeline_output import StableDiffusionXLPipelineOutput from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import rescale_noise_cfg from diffusers.utils import is_torch_xla_available -from k_diffusion.external import CompVisVDenoiser, CompVisDenoiser -from k_diffusion.sampling import get_sigmas_karras, BrownianTreeNoiseSampler from toolkit.models.flux import bypass_flux_guidance, restore_flux_guidance from diffusers.image_processor import PipelineImageInput from PIL import Image @@ -29,278 +24,6 @@ if is_torch_xla_available(): else: XLA_AVAILABLE = False -class StableDiffusionKDiffusionXLPipeline(StableDiffusionXLPipeline): - - def __init__( - self, - vae: 'AutoencoderKL', - text_encoder: 'CLIPTextModel', - text_encoder_2: 'CLIPTextModelWithProjection', - tokenizer: 'CLIPTokenizer', - tokenizer_2: 'CLIPTokenizer', - unet: 'UNet2DConditionModel', - scheduler: 'KarrasDiffusionSchedulers', - force_zeros_for_empty_prompt: bool = True, - add_watermarker: Optional[bool] = None, - ): - super().__init__( - vae=vae, - text_encoder=text_encoder, - text_encoder_2=text_encoder_2, - tokenizer=tokenizer, - tokenizer_2=tokenizer_2, - unet=unet, - scheduler=scheduler, - ) - raise NotImplementedError("This pipeline is not implemented yet") - # self.sampler = None - # scheduler = LMSDiscreteScheduler.from_config(scheduler.config) - # model = ModelWrapper(unet, scheduler.alphas_cumprod) - # if scheduler.config.prediction_type == "v_prediction": - # self.k_diffusion_model = CompVisVDenoiser(model) - # else: - # self.k_diffusion_model = CompVisDenoiser(model) - - def set_scheduler(self, scheduler_type: str): - library = importlib.import_module("k_diffusion") - sampling = getattr(library, "sampling") - self.sampler = getattr(sampling, scheduler_type) - - @torch.no_grad() - def __call__( - self, - prompt: Union[str, List[str]] = None, - prompt_2: Optional[Union[str, List[str]]] = None, - height: Optional[int] = None, - width: Optional[int] = None, - num_inference_steps: int = 50, - denoising_end: Optional[float] = None, - guidance_scale: float = 5.0, - negative_prompt: Optional[Union[str, List[str]]] = None, - negative_prompt_2: Optional[Union[str, List[str]]] = None, - num_images_per_prompt: Optional[int] = 1, - eta: float = 0.0, - generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, - latents: Optional[torch.FloatTensor] = None, - prompt_embeds: Optional[torch.FloatTensor] = None, - negative_prompt_embeds: Optional[torch.FloatTensor] = None, - pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None, - output_type: Optional[str] = "pil", - return_dict: bool = True, - callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None, - callback_steps: int = 1, - cross_attention_kwargs: Optional[Dict[str, Any]] = None, - guidance_rescale: float = 0.0, - original_size: Optional[Tuple[int, int]] = None, - crops_coords_top_left: Tuple[int, int] = (0, 0), - target_size: Optional[Tuple[int, int]] = None, - use_karras_sigmas: bool = False, - ): - - # 0. Default height and width to unet - height = height or self.default_sample_size * self.vae_scale_factor - width = width or self.default_sample_size * self.vae_scale_factor - - original_size = original_size or (height, width) - target_size = target_size or (height, width) - - # 1. Check inputs. Raise error if not correct - self.check_inputs( - prompt, - prompt_2, - height, - width, - callback_steps, - negative_prompt, - negative_prompt_2, - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) - - # 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 = self._execution_device - - # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2) - # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1` - # corresponds to doing no classifier free guidance. - do_classifier_free_guidance = guidance_scale > 1.0 - - # 3. Encode input prompt - text_encoder_lora_scale = ( - cross_attention_kwargs.get("scale", None) if cross_attention_kwargs is not None else None - ) - ( - prompt_embeds, - negative_prompt_embeds, - pooled_prompt_embeds, - negative_pooled_prompt_embeds, - ) = self.encode_prompt( - prompt=prompt, - prompt_2=prompt_2, - device=device, - num_images_per_prompt=num_images_per_prompt, - do_classifier_free_guidance=do_classifier_free_guidance, - negative_prompt=negative_prompt, - negative_prompt_2=negative_prompt_2, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - pooled_prompt_embeds=pooled_prompt_embeds, - negative_pooled_prompt_embeds=negative_pooled_prompt_embeds, - lora_scale=text_encoder_lora_scale, - ) - - # 4. Prepare timesteps - self.scheduler.set_timesteps(num_inference_steps, device=device) - - timesteps = self.scheduler.timesteps - - # 5. Prepare latent variables - 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, - generator, - latents, - ) - - # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline - extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta) - - # 7. Prepare added time ids & embeddings - add_text_embeds = pooled_prompt_embeds - add_time_ids = self._get_add_time_ids( - original_size, crops_coords_top_left, target_size, dtype=prompt_embeds.dtype - ) - - if do_classifier_free_guidance: - prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) - add_text_embeds = torch.cat([negative_pooled_prompt_embeds, add_text_embeds], dim=0) - add_time_ids = torch.cat([add_time_ids, add_time_ids], dim=0) - - prompt_embeds = prompt_embeds.to(device) - add_text_embeds = add_text_embeds.to(device) - add_time_ids = add_time_ids.to(device).repeat(batch_size * num_images_per_prompt, 1) - - # 8. Denoising loop - num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) - - # 7.1 Apply denoising_end - if denoising_end is not None and type(denoising_end) == float and denoising_end > 0 and denoising_end < 1: - discrete_timestep_cutoff = int( - round( - self.scheduler.config.num_train_timesteps - - (denoising_end * self.scheduler.config.num_train_timesteps) - ) - ) - num_inference_steps = len(list(filter(lambda ts: ts >= discrete_timestep_cutoff, timesteps))) - timesteps = timesteps[:num_inference_steps] - - # 5. Prepare sigmas - if use_karras_sigmas: - sigma_min: float = self.k_diffusion_model.sigmas[0].item() - sigma_max: float = self.k_diffusion_model.sigmas[-1].item() - sigmas = get_sigmas_karras(n=num_inference_steps, sigma_min=sigma_min, sigma_max=sigma_max) - sigmas = sigmas.to(device) - else: - sigmas = self.scheduler.sigmas - sigmas = sigmas.to(prompt_embeds.dtype) - - # 5. Prepare latent variables - 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, - generator, - latents, - ) - - latents = latents * sigmas[0] - self.k_diffusion_model.sigmas = self.k_diffusion_model.sigmas.to(latents.device) - self.k_diffusion_model.log_sigmas = self.k_diffusion_model.log_sigmas.to(latents.device) - - # 7. Define model function - def model_fn(x, t): - latent_model_input = torch.cat([x] * 2) - t = torch.cat([t] * 2) - - added_cond_kwargs = {"text_embeds": add_text_embeds, "time_ids": add_time_ids} - # noise_pred = self.unet( - # latent_model_input, - # t, - # encoder_hidden_states=prompt_embeds, - # cross_attention_kwargs=cross_attention_kwargs, - # added_cond_kwargs=added_cond_kwargs, - # return_dict=False, - # )[0] - - noise_pred = self.k_diffusion_model( - latent_model_input, - t, - encoder_hidden_states=prompt_embeds, - cross_attention_kwargs=cross_attention_kwargs, - added_cond_kwargs=added_cond_kwargs, - return_dict=False,)[0] - - noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) - noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) - return noise_pred - - - # 8. Run k-diffusion solver - sampler_kwargs = {} - # should work without it - noise_sampler_seed = None - - - if "noise_sampler" in inspect.signature(self.sampler).parameters: - min_sigma, max_sigma = sigmas[sigmas > 0].min(), sigmas.max() - noise_sampler = BrownianTreeNoiseSampler(latents, min_sigma, max_sigma, noise_sampler_seed) - sampler_kwargs["noise_sampler"] = noise_sampler - - latents = self.sampler(model_fn, latents, sigmas, **sampler_kwargs) - - if not output_type == "latent": - image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0] - image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype) - else: - image = latents - 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) - - # Offload last model to CPU - if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None: - self.final_offload_hook.offload() - - if not return_dict: - return (image,) - - return StableDiffusionXLPipelineOutput(images=image) - - class CustomStableDiffusionXLPipeline(StableDiffusionXLPipeline): def predict_noise( diff --git a/toolkit/sampler.py b/toolkit/sampler.py index c36af3d..daed568 100644 --- a/toolkit/sampler.py +++ b/toolkit/sampler.py @@ -20,8 +20,6 @@ from toolkit.samplers.mean_flow_scheduler import MeanFlowScheduler from toolkit.samplers.custom_flowmatch_sampler import CustomFlowMatchEulerDiscreteScheduler -from k_diffusion.external import CompVisDenoiser - from toolkit.samplers.custom_lcm_scheduler import CustomLCMScheduler # scheduler: @@ -185,29 +183,3 @@ def get_sampler( scheduler = scheduler_cls.from_config(config) return scheduler - - -# testing -if __name__ == "__main__": - from diffusers import DiffusionPipeline - - from diffusers import StableDiffusionKDiffusionPipeline - import torch - import os - - inference_steps = 25 - - pipe = StableDiffusionKDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1-base") - pipe = pipe.to("cuda") - - k_diffusion_model = CompVisDenoiser(model) - - pipe = DiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", custom_pipeline="sd_text2img_k_diffusion") - pipe = pipe.to("cuda") - - prompt = "an astronaut riding a horse on mars" - pipe.set_scheduler("sample_heun") - generator = torch.Generator(device="cuda").manual_seed(seed) - image = pipe(prompt, generator=generator, num_inference_steps=20).images[0] - - image.save("./astronaut_heun_k_diffusion.png") diff --git a/toolkit/stable_diffusion_model.py b/toolkit/stable_diffusion_model.py index f574deb..915bc00 100644 --- a/toolkit/stable_diffusion_model.py +++ b/toolkit/stable_diffusion_model.py @@ -42,7 +42,7 @@ from toolkit.train_tools import get_torch_dtype, apply_noise_offset from einops import rearrange, repeat import torch from toolkit.pipelines import CustomStableDiffusionXLPipeline, CustomStableDiffusionPipeline, \ - StableDiffusionKDiffusionXLPipeline, StableDiffusionXLRefinerPipeline, FluxWithCFGPipeline, \ + StableDiffusionXLRefinerPipeline, FluxWithCFGPipeline, \ FluxAdvancedControlPipeline from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, T2IAdapter, DDPMScheduler, \ StableDiffusionXLAdapterPipeline, StableDiffusionAdapterPipeline, DiffusionPipeline, PixArtTransformer2DModel, \ @@ -1212,10 +1212,7 @@ class StableDiffusion: except: pass - if sampler.startswith("sample_") and self.is_xl: - # using kdiffusion - Pipe = StableDiffusionKDiffusionXLPipeline - elif self.is_xl: + if self.is_xl: Pipe = StableDiffusionXLPipeline elif self.is_v3: Pipe = StableDiffusion3Pipeline @@ -1348,9 +1345,6 @@ class StableDiffusion: # disable progress bar pipeline.set_progress_bar_config(disable=True) - if sampler.startswith("sample_"): - pipeline.set_scheduler(sampler) - refiner_pipeline = None if self.refiner_unet: # build refiner pipeline