Big refactor of SD runner and added image generator

toolkit/config_modules.py

@@ -1,4 +1,7 @@
-from typing import List
+import os
+import time
+from typing import List, Optional
+import random
 
 
 class SaveConfig:
@@ -27,6 +30,7 @@ class SampleConfig:
         self.guidance_scale = kwargs.get('guidance_scale', 7)
         self.sample_steps = kwargs.get('sample_steps', 20)
         self.network_multiplier = kwargs.get('network_multiplier', 1)
+        self.guidance_rescale = kwargs.get('guidance_rescale', 0.0)
 
 
 class NetworkConfig:
@@ -35,7 +39,7 @@ class NetworkConfig:
         rank = kwargs.get('rank', None)
         linear = kwargs.get('linear', None)
         if rank is not None:
-            self.rank: int = rank # rank for backward compatibility
+            self.rank: int = rank  # rank for backward compatibility
             self.linear: int = rank
         elif linear is not None:
             self.rank: int = linear
@@ -71,6 +75,7 @@ class ModelConfig:
         self.is_v2: bool = kwargs.get('is_v2', False)
         self.is_xl: bool = kwargs.get('is_xl', False)
         self.is_v_pred: bool = kwargs.get('is_v_pred', False)
+        self.dtype: str = kwargs.get('dtype', 'float16')
 
         if self.name_or_path is None:
             raise ValueError('name_or_path must be specified')
@@ -103,3 +108,197 @@ class SliderConfig:
         self.resolutions: List[List[int]] = kwargs.get('resolutions', [[512, 512]])
         self.prompt_file: str = kwargs.get('prompt_file', None)
         self.prompt_tensors: str = kwargs.get('prompt_tensors', None)
+
+
+class GenerateImageConfig:
+    def __init__(
+            self,
+            prompt: str = '',
+            prompt_2: Optional[str] = None,
+            width: int = 512,
+            height: int = 512,
+            num_inference_steps: int = 50,
+            guidance_scale: float = 7.5,
+            negative_prompt: str = '',
+            negative_prompt_2: Optional[str] = None,
+            seed: int = -1,
+            network_multiplier: float = 1.0,
+            guidance_rescale: float = 0.0,
+            # the tag [time] will be replaced with milliseconds since epoch
+            output_path: str = None,  # full image path
+            output_folder: str = None,  # folder to save image in if output_path is not specified
+            output_ext: str = 'png',  # extension to save image as if output_path is not specified
+            output_tail: str = '',  # tail to add to output filename
+            add_prompt_file: bool = False,  # add a prompt file with generated image
+    ):
+        self.width: int = width
+        self.height: int = height
+        self.num_inference_steps: int = num_inference_steps
+        self.guidance_scale: float = guidance_scale
+        self.guidance_rescale: float = guidance_rescale
+        self.prompt: str = prompt
+        self.prompt_2: str = prompt_2
+        self.negative_prompt: str = negative_prompt
+        self.negative_prompt_2: str = negative_prompt_2
+
+        self.output_path: str = output_path
+        self.seed: int = seed
+        if self.seed == -1:
+            # generate random one
+            self.seed = random.randint(0, 2 ** 32 - 1)
+        self.network_multiplier: float = network_multiplier
+        self.output_folder: str = output_folder
+        self.output_ext: str = output_ext
+        self.add_prompt_file: bool = add_prompt_file
+        self.output_tail: str = output_tail
+        self.gen_time: int = int(time.time() * 1000)
+
+        # prompt string will override any settings above
+        self._process_prompt_string()
+
+        # handle dual text encoder prompts if nothing passed
+        if negative_prompt_2 is None:
+            self.negative_prompt_2 = negative_prompt
+
+        if prompt_2 is None:
+            self.prompt_2 = prompt
+
+        # parse prompt paths
+        if self.output_path is None and self.output_folder is None:
+            raise ValueError('output_path or output_folder must be specified')
+        elif self.output_path is not None:
+            self.output_folder = os.path.dirname(self.output_path)
+            self.output_ext = os.path.splitext(self.output_path)[1][1:]
+            self.output_filename_no_ext = os.path.splitext(os.path.basename(self.output_path))[0]
+
+        else:
+            self.output_filename_no_ext = '[time]_[count]'
+            if len(self.output_tail) > 0:
+                self.output_filename_no_ext += '_' + self.output_tail
+            self.output_path = os.path.join(self.output_folder, self.output_filename_no_ext + '.' + self.output_ext)
+
+        # adjust height
+        self.height = max(64, self.height - self.height % 8)  # round to divisible by 8
+        self.width = max(64, self.width - self.width % 8)  # round to divisible by 8
+
+    def set_gen_time(self, gen_time: int = None):
+        if gen_time is not None:
+            self.gen_time = gen_time
+        else:
+            self.gen_time = int(time.time() * 1000)
+
+    def _get_path_no_ext(self, count: int = 0, max_count=0):
+        # zero pad count
+        count_str = str(count).zfill(len(str(max_count)))
+        # replace [time] with gen time
+        filename = self.output_filename_no_ext.replace('[time]', str(self.gen_time))
+        # replace [count] with count
+        filename = filename.replace('[count]', count_str)
+        return filename
+
+    def get_image_path(self, count: int = 0, max_count=0):
+        filename = self._get_path_no_ext(count, max_count)
+        filename += '.' + self.output_ext
+        # join with folder
+        return os.path.join(self.output_folder, filename)
+
+    def get_prompt_path(self, count: int = 0, max_count=0):
+        filename = self._get_path_no_ext(count, max_count)
+        filename += '.txt'
+        # join with folder
+        return os.path.join(self.output_folder, filename)
+
+    def save_image(self, image, count: int = 0, max_count=0):
+        # make parent dirs
+        os.makedirs(self.output_folder, exist_ok=True)
+        self.set_gen_time()
+        # TODO save image gen header info for A1111 and us, our seeds probably wont match
+        image.save(self.get_image_path(count, max_count))
+        # do prompt file
+        if self.add_prompt_file:
+            self.save_prompt_file(count, max_count)
+
+    def save_prompt_file(self, count: int = 0, max_count=0):
+        # save prompt file
+        with open(self.get_prompt_path(count, max_count), 'w') as f:
+            prompt = self.prompt
+            if self.prompt_2 is not None:
+                prompt += ' --p2 ' + self.prompt_2
+            if self.negative_prompt is not None:
+                prompt += ' --n ' + self.negative_prompt
+            if self.negative_prompt_2 is not None:
+                prompt += ' --n2 ' + self.negative_prompt_2
+            prompt += ' --w ' + str(self.width)
+            prompt += ' --h ' + str(self.height)
+            prompt += ' --seed ' + str(self.seed)
+            prompt += ' --cfg ' + str(self.guidance_scale)
+            prompt += ' --steps ' + str(self.num_inference_steps)
+            prompt += ' --m ' + str(self.network_multiplier)
+            prompt += ' --gr ' + str(self.guidance_rescale)
+
+            # get gen info
+            f.write(self.prompt)
+
+    def _process_prompt_string(self):
+        # we will try to support all sd-scripts where we can
+
+        # FROM SD-SCRIPTS
+        # --n Treat everything until the next option as a negative prompt.
+        # --w Specify the width of the generated image.
+        # --h Specify the height of the generated image.
+        # --d Specify the seed for the generated image.
+        # --l Specify the CFG scale for the generated image.
+        # --s Specify the number of steps during generation.
+
+        # OURS and some QOL additions
+        # --m Specify the network multiplier for the generated image.
+        # --p2 Prompt for the second text encoder (SDXL only)
+        # --n2 Negative prompt for the second text encoder (SDXL only)
+        # --gr Specify the guidance rescale for the generated image (SDXL only)
+
+        # --seed Specify the seed for the generated image same as --d
+        # --cfg Specify the CFG scale for the generated image same as --l
+        # --steps Specify the number of steps during generation same as --s
+        # --network_multiplier Specify the network multiplier for the generated image same as --m
+
+        # process prompt string and update values if it has some
+        if self.prompt is not None and len(self.prompt) > 0:
+            # process prompt string
+            prompt = self.prompt
+            prompt = prompt.strip()
+            p_split = prompt.split('--')
+            self.prompt = p_split[0].strip()
+
+            if len(p_split) > 1:
+                for split in p_split[1:]:
+                    # allows multi char flags
+                    flag = split.split(' ')[0].strip()
+                    content = split[len(flag):].strip()
+                    if flag == 'p2':
+                        self.prompt_2 = content
+                    elif flag == 'n':
+                        self.negative_prompt = content
+                    elif flag == 'n2':
+                        self.negative_prompt_2 = content
+                    elif flag == 'w':
+                        self.width = int(content)
+                    elif flag == 'h':
+                        self.height = int(content)
+                    elif flag == 'd':
+                        self.seed = int(content)
+                    elif flag == 'seed':
+                        self.seed = int(content)
+                    elif flag == 'l':
+                        self.guidance_scale = float(content)
+                    elif flag == 'cfg':
+                        self.guidance_scale = float(content)
+                    elif flag == 's':
+                        self.num_inference_steps = int(content)
+                    elif flag == 'steps':
+                        self.num_inference_steps = int(content)
+                    elif flag == 'm':
+                        self.network_multiplier = float(content)
+                    elif flag == 'network_multiplier':
+                        self.network_multiplier = float(content)
+                    elif flag == 'gr':
+                        self.guidance_rescale = float(content)

toolkit/job.py

@@ -16,6 +16,9 @@ def get_job(config_path, name=None):
     if job == 'mod':
         from jobs import ModJob
         return ModJob(config)
+    if job == 'generate':
+        from jobs import GenerateJob
+        return GenerateJob(config)
 
     # elif job == 'train':
     #     from jobs import TrainJob

toolkit/scheduler.py

@@ -0,0 +1,33 @@
+import torch
+from typing import Optional
+
+
+def get_lr_scheduler(
+        name: Optional[str],
+        optimizer: torch.optim.Optimizer,
+        max_iterations: Optional[int],
+        lr_min: Optional[float],
+        **kwargs,
+):
+    if name == "cosine":
+        return torch.optim.lr_scheduler.CosineAnnealingLR(
+            optimizer, T_max=max_iterations, eta_min=lr_min, **kwargs
+        )
+    elif name == "cosine_with_restarts":
+        return torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
+            optimizer, T_0=max_iterations, T_mult=2, eta_min=lr_min, **kwargs
+        )
+    elif name == "step":
+        return torch.optim.lr_scheduler.StepLR(
+            optimizer, step_size=max_iterations // 100, gamma=0.999, **kwargs
+        )
+    elif name == "constant":
+        return torch.optim.lr_scheduler.ConstantLR(optimizer, factor=1, **kwargs)
+    elif name == "linear":
+        return torch.optim.lr_scheduler.LinearLR(
+            optimizer, start_factor=0.5, end_factor=0.5, total_iters=max_iterations, **kwargs
+        )
+    else:
+        raise ValueError(
+            "Scheduler must be cosine, cosine_with_restarts, step, linear or constant"
+        )

toolkit/stable_diffusion_model.py

@@ -1,12 +1,16 @@
+import gc
 import typing
-from typing import Union, OrderedDict
+from typing import Union, OrderedDict, List
 import sys
 import os
 
+from diffusers.pipelines.stable_diffusion_xl.pipeline_stable_diffusion_xl import rescale_noise_cfg
 from safetensors.torch import save_file
+from tqdm import tqdm
 
+from toolkit.config_modules import ModelConfig, GenerateImageConfig
 from toolkit.paths import REPOS_ROOT
-from toolkit.train_tools import get_torch_dtype
+from toolkit.train_tools import get_torch_dtype, apply_noise_offset
 
 sys.path.append(REPOS_ROOT)
 sys.path.append(os.path.join(REPOS_ROOT, 'leco'))
@@ -14,6 +18,32 @@ from leco import train_util
 import torch
 from library import model_util
 from library.sdxl_model_util import convert_text_encoder_2_state_dict_to_sdxl
+from diffusers.schedulers import DDPMScheduler
+from toolkit.pipelines import CustomStableDiffusionXLPipeline, CustomStableDiffusionPipeline
+from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
+
+
+class BlankNetwork:
+    multiplier = 1.0
+    is_active = True
+
+    def __init__(self):
+        pass
+
+    def __enter__(self):
+        self.is_active = True
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.is_active = False
+
+
+def flush():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+
+UNET_IN_CHANNELS = 4  # Stable Diffusion の in_channels は 4 で固定。XLも同じ。
+VAE_SCALE_FACTOR = 8  # 2 ** (len(vae.config.block_out_channels) - 1) = 8
 
 
 class PromptEmbeds:
@@ -39,31 +69,382 @@ class PromptEmbeds:
 
 # if is type checking
 if typing.TYPE_CHECKING:
-    from diffusers import StableDiffusionPipeline
-    from toolkit.pipelines import CustomStableDiffusionXLPipeline
+    from diffusers import \
+        StableDiffusionPipeline, \
+        AutoencoderKL, \
+        UNet2DConditionModel
+    from diffusers.schedulers import KarrasDiffusionSchedulers
+    from transformers import CLIPTextModel, CLIPTokenizer, CLIPTextModelWithProjection
 
 
 class StableDiffusion:
     pipeline: Union[None, 'StableDiffusionPipeline', 'CustomStableDiffusionXLPipeline']
+    vae: Union[None, 'AutoencoderKL']
+    unet: Union[None, 'UNet2DConditionModel']
+    text_encoder: Union[None, 'CLIPTextModel', List[Union['CLIPTextModel', 'CLIPTextModelWithProjection']]]
+    tokenizer: Union[None, 'CLIPTokenizer', List['CLIPTokenizer']]
+    noise_scheduler: Union[None, 'KarrasDiffusionSchedulers', 'DDPMScheduler']
+    device: str
+    dtype: str
+    torch_dtype: torch.dtype
+    device_torch: torch.device
+    model_config: ModelConfig
 
     def __init__(
             self,
-            vae,
-            tokenizer,
-            text_encoder,
-            unet,
-            noise_scheduler,
-            is_xl=False,
-            pipeline=None,
+            device,
+            model_config: ModelConfig,
+            dtype='fp16',
+            custom_pipeline=None
     ):
-        # text encoder has a list of 2 for xl
-        self.vae = vae
+        self.custom_pipeline = custom_pipeline
+        self.device = device
+        self.dtype = dtype
+        self.torch_dtype = get_torch_dtype(dtype)
+        self.device_torch = torch.device(self.device)
+        self.model_config = model_config
+        self.prediction_type = "v_prediction" if self.model_config.is_v_pred else "epsilon"
+
+        # sdxl stuff
+        self.logit_scale = None
+        self.ckppt_info = None
+
+        # to hold network if there is one
+        self.network = None
+        self.is_xl = model_config.is_xl
+        self.is_v2 = model_config.is_v2
+
+    def load_model(self):
+        dtype = get_torch_dtype(self.dtype)
+
+        # TODO handle other schedulers
+        # sch = KDPM2DiscreteScheduler
+        sch = DDPMScheduler
+        # do our own scheduler
+        prediction_type = "v_prediction" if self.model_config.is_v_pred else "epsilon"
+        scheduler = sch(
+            num_train_timesteps=1000,
+            beta_start=0.00085,
+            beta_end=0.0120,
+            beta_schedule="scaled_linear",
+            clip_sample=False,
+            prediction_type=prediction_type,
+            steps_offset=1
+        )
+        if self.model_config.is_xl:
+            if self.custom_pipeline is not None:
+                pipln = self.custom_pipeline
+            else:
+                pipln = CustomStableDiffusionXLPipeline
+            pipe = pipln.from_single_file(
+                self.model_config.name_or_path,
+                dtype=dtype,
+                scheduler_type='ddpm',
+                device=self.device_torch,
+            ).to(self.device_torch)
+
+            text_encoders = [pipe.text_encoder, pipe.text_encoder_2]
+            tokenizer = [pipe.tokenizer, pipe.tokenizer_2]
+            for text_encoder in text_encoders:
+                text_encoder.to(self.device_torch, dtype=dtype)
+                text_encoder.requires_grad_(False)
+                text_encoder.eval()
+            text_encoder = text_encoders
+        else:
+            if self.custom_pipeline is not None:
+                pipln = self.custom_pipeline
+            else:
+                pipln = CustomStableDiffusionPipeline
+            pipe = pipln.from_single_file(
+                self.model_config.name_or_path,
+                dtype=dtype,
+                scheduler_type='dpm',
+                device=self.device_torch,
+                load_safety_checker=False,
+            ).to(self.device_torch)
+            pipe.register_to_config(requires_safety_checker=False)
+            text_encoder = pipe.text_encoder
+            text_encoder.to(self.device_torch, dtype=dtype)
+            text_encoder.requires_grad_(False)
+            text_encoder.eval()
+            tokenizer = pipe.tokenizer
+
+        # scheduler doesn't get set sometimes, so we set it here
+        pipe.scheduler = scheduler
+
+        self.unet = pipe.unet
+        self.noise_scheduler = pipe.scheduler
+        self.vae = pipe.vae.to(self.device_torch, dtype=dtype)
+        self.vae.eval()
+        self.vae.requires_grad_(False)
+        self.unet.to(self.device_torch, dtype=dtype)
+        self.unet.requires_grad_(False)
+        self.unet.eval()
+
         self.tokenizer = tokenizer
         self.text_encoder = text_encoder
-        self.unet = unet
-        self.noise_scheduler = noise_scheduler
-        self.is_xl = is_xl
-        self.pipeline = pipeline
+        self.pipeline = pipe
+
+    def generate_images(self, image_configs: List[GenerateImageConfig]):
+        # sample_folder = os.path.join(self.save_root, 'samples')
+        if self.network is not None:
+            self.network.eval()
+            network = self.network
+        else:
+            network = BlankNetwork()
+
+        # save current seed state for training
+        rng_state = torch.get_rng_state()
+        cuda_rng_state = torch.cuda.get_rng_state() if torch.cuda.is_available() else None
+
+        original_device_dict = {
+            'vae': self.vae.device,
+            'unet': self.unet.device,
+            # 'tokenizer': self.tokenizer.device,
+        }
+
+        # handle sdxl text encoder
+        if isinstance(self.text_encoder, list):
+            for encoder, i in zip(self.text_encoder, range(len(self.text_encoder))):
+                original_device_dict[f'text_encoder_{i}'] = encoder.device
+                encoder.to(self.device_torch)
+        else:
+            original_device_dict['text_encoder'] = self.text_encoder.device
+            self.text_encoder.to(self.device_torch)
+
+        self.vae.to(self.device_torch)
+        self.unet.to(self.device_torch)
+
+        # TODO add clip skip
+        if self.is_xl:
+            pipeline = StableDiffusionXLPipeline(
+                vae=self.vae,
+                unet=self.unet,
+                text_encoder=self.text_encoder[0],
+                text_encoder_2=self.text_encoder[1],
+                tokenizer=self.tokenizer[0],
+                tokenizer_2=self.tokenizer[1],
+                scheduler=self.noise_scheduler,
+                add_watermarker=False,
+            ).to(self.device_torch)
+            # force turn that (ruin your images with obvious green and red dots) the #$@@ off!!!
+            pipeline.watermark = None
+        else:
+            pipeline = StableDiffusionPipeline(
+                vae=self.vae,
+                unet=self.unet,
+                text_encoder=self.text_encoder,
+                tokenizer=self.tokenizer,
+                scheduler=self.noise_scheduler,
+                safety_checker=None,
+                feature_extractor=None,
+                requires_safety_checker=False,
+            ).to(self.device_torch)
+        # disable progress bar
+        pipeline.set_progress_bar_config(disable=True)
+
+        start_multiplier = 1.0
+        if self.network is not None:
+            start_multiplier = self.network.multiplier
+
+        pipeline.to(self.device_torch)
+        with network:
+            with torch.no_grad():
+                if self.network is not None:
+                    assert self.network.is_active
+
+                for i in tqdm(range(len(image_configs)), desc=f"Generating Images", leave=False):
+                    gen_config = image_configs[i]
+
+                    if self.network is not None:
+                        self.network.multiplier = gen_config.network_multiplier
+                    torch.manual_seed(gen_config.seed)
+                    torch.cuda.manual_seed(gen_config.seed)
+
+                    if self.is_xl:
+                        img = pipeline(
+                            prompt=gen_config.prompt,
+                            prompt_2=gen_config.prompt_2,
+                            negative_prompt=gen_config.negative_prompt,
+                            negative_prompt_2=gen_config.negative_prompt_2,
+                            height=gen_config.height,
+                            width=gen_config.width,
+                            num_inference_steps=gen_config.num_inference_steps,
+                            guidance_scale=gen_config.guidance_scale,
+                            guidance_rescale=gen_config.guidance_rescale,
+                        ).images[0]
+                    else:
+                        img = pipeline(
+                            prompt=gen_config.prompt,
+                            negative_prompt=gen_config.negative_prompt,
+                            height=gen_config.height,
+                            width=gen_config.width,
+                            num_inference_steps=gen_config.num_inference_steps,
+                            guidance_scale=gen_config.guidance_scale,
+                        ).images[0]
+
+                    gen_config.save_image(img)
+
+        # clear pipeline and cache to reduce vram usage
+        del pipeline
+        torch.cuda.empty_cache()
+
+        # restore training state
+        torch.set_rng_state(rng_state)
+        if cuda_rng_state is not None:
+            torch.cuda.set_rng_state(cuda_rng_state)
+
+        self.vae.to(original_device_dict['vae'])
+        self.unet.to(original_device_dict['unet'])
+        if isinstance(self.text_encoder, list):
+            for encoder, i in zip(self.text_encoder, range(len(self.text_encoder))):
+                encoder.to(original_device_dict[f'text_encoder_{i}'])
+        else:
+            self.text_encoder.to(original_device_dict['text_encoder'])
+        if self.network is not None:
+            self.network.train()
+            self.network.multiplier = start_multiplier
+        # self.tokenizer.to(original_device_dict['tokenizer'])
+
+    def get_latent_noise(
+            self,
+            height=None,
+            width=None,
+            pixel_height=None,
+            pixel_width=None,
+            batch_size=1,
+            noise_offset=0.0,
+    ):
+        if height is None and pixel_height is None:
+            raise ValueError("height or pixel_height must be specified")
+        if width is None and pixel_width is None:
+            raise ValueError("width or pixel_width must be specified")
+        if height is None:
+            height = pixel_height // VAE_SCALE_FACTOR
+        if width is None:
+            width = pixel_width // VAE_SCALE_FACTOR
+
+        noise = torch.randn(
+            (
+                batch_size,
+                UNET_IN_CHANNELS,
+                height,
+                width,
+            ),
+            device="cpu",
+        )
+        noise = apply_noise_offset(noise, noise_offset)
+        return noise
+
+    def get_time_ids_from_latents(self, latents: torch.Tensor):
+        bs, ch, h, w = list(latents.shape)
+
+        height = h * VAE_SCALE_FACTOR
+        width = w * VAE_SCALE_FACTOR
+
+        dtype = latents.dtype
+
+        if self.is_xl:
+            prompt_ids = train_util.get_add_time_ids(
+                height,
+                width,
+                dynamic_crops=False,  # look into this
+                dtype=dtype,
+            ).to(self.device_torch, dtype=dtype)
+            return train_util.concat_embeddings(
+                prompt_ids, prompt_ids, bs
+            )
+        else:
+            return None
+
+    def predict_noise(
+            self,
+            latents: torch.FloatTensor,
+            text_embeddings: PromptEmbeds,
+            timestep: int,
+            guidance_scale=7.5,
+            guidance_rescale=0,  # 0.7
+            add_time_ids=None,
+            **kwargs,
+    ):
+
+        if self.is_xl:
+            if add_time_ids is None:
+                add_time_ids = self.get_time_ids_from_latents(latents)
+
+            latent_model_input = torch.cat([latents] * 2)
+
+            latent_model_input = self.noise_scheduler.scale_model_input(latent_model_input, timestep)
+
+            added_cond_kwargs = {
+                "text_embeds": text_embeddings.pooled_embeds,
+                "time_ids": add_time_ids,
+            }
+
+            # predict the noise residual
+            noise_pred = self.unet(
+                latent_model_input,
+                timestep,
+                encoder_hidden_states=text_embeddings.text_embeds,
+                added_cond_kwargs=added_cond_kwargs,
+            ).sample
+
+            # perform guidance
+            noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (
+                    noise_pred_text - noise_pred_uncond
+            )
+
+            # https://github.com/huggingface/diffusers/blob/7a91ea6c2b53f94da930a61ed571364022b21044/src/diffusers/pipelines/stable_diffusion_xl/pipeline_stable_diffusion_xl.py#L775
+            if guidance_rescale > 0.0:
+                # Based on 3.4. in https://arxiv.org/pdf/2305.08891.pdf
+                noise_pred = rescale_noise_cfg(noise_pred, noise_pred_text, guidance_rescale=guidance_rescale)
+
+        else:
+            # if we are doing classifier free guidance, need to double up
+            latent_model_input = torch.cat([latents] * 2)
+
+            latent_model_input = self.noise_scheduler.scale_model_input(latent_model_input, timestep)
+
+            # predict the noise residual
+            noise_pred = self.unet(
+                latent_model_input,
+                timestep,
+                encoder_hidden_states=text_embeddings.text_embeds,
+            ).sample
+
+            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
+
+    # ref: https://github.com/huggingface/diffusers/blob/0bab447670f47c28df60fbd2f6a0f833f75a16f5/src/diffusers/pipelines/stable_diffusion/pipeline_stable_diffusion.py#L746
+    def diffuse_some_steps(
+            self,
+            latents: torch.FloatTensor,
+            text_embeddings: PromptEmbeds,
+            total_timesteps: int = 1000,
+            start_timesteps=0,
+            guidance_scale=1,
+            add_time_ids=None,
+            **kwargs,
+    ):
+
+        for timestep in tqdm(self.noise_scheduler.timesteps[start_timesteps:total_timesteps], leave=False):
+            noise_pred = self.predict_noise(
+                latents,
+                text_embeddings,
+                timestep,
+                guidance_scale=guidance_scale,
+                add_time_ids=add_time_ids,
+                **kwargs,
+            )
+            latents = self.noise_scheduler.step(noise_pred, timestep, latents).prev_sample
+
+        # return latents_steps
+        return latents
 
     def encode_prompt(self, prompt, num_images_per_prompt=1) -> PromptEmbeds:
         prompt = prompt