diff --git a/pipeline_fill_sd_xl.py b/pipeline_fill_sd_xl.py index 8e24bd7..297f74a 100644 --- a/pipeline_fill_sd_xl.py +++ b/pipeline_fill_sd_xl.py @@ -17,12 +17,10 @@ from typing import List, Optional, Union import cv2 import PIL.Image import torch -import gc from diffusers.image_processor import PipelineImageInput, VaeImageProcessor -from diffusers.models import AutoencoderKL, UNet2DConditionModel -from diffusers.pipelines.pipeline_utils import DiffusionPipeline, StableDiffusionMixin from diffusers.schedulers import KarrasDiffusionSchedulers from diffusers.utils.torch_utils import randn_tensor +from tqdm import tqdm from .controlnet_union import ControlNetModel_Union from comfy.utils import ProgressBar @@ -71,17 +69,9 @@ class StableDiffusionXLFillPipeline(DiffusionPipeline, StableDiffusionMixin): def __init__( self, - unet: UNet2DConditionModel, - scheduler: KarrasDiffusionSchedulers, - force_zeros_for_empty_prompt: bool = True, ): super().__init__() - self.register_modules( - unet=unet, - scheduler=scheduler, - ) - self.vae_scale_factor = 8 self.image_processor = VaeImageProcessor( vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True @@ -91,10 +81,6 @@ class StableDiffusionXLFillPipeline(DiffusionPipeline, StableDiffusionMixin): do_convert_rgb=True, do_normalize=False, ) - self.register_to_config( - force_zeros_for_empty_prompt=force_zeros_for_empty_prompt - ) - self.controlnet_model = None def prepare_image(self, image, device, dtype, do_classifier_free_guidance=False): image = self.control_image_processor.preprocess(image).to(dtype=torch.float32) @@ -134,7 +120,10 @@ class StableDiffusionXLFillPipeline(DiffusionPipeline, StableDiffusionMixin): # corresponds to doing no classifier free guidance. @property def do_classifier_free_guidance(self): - return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None + if hasattr(self.unet, 'config'): + return self._guidance_scale > 1 and self.unet.config.time_cond_proj_dim is None + else: + return self._guidance_scale > 1 @property def num_timesteps(self): @@ -147,6 +136,10 @@ class StableDiffusionXLFillPipeline(DiffusionPipeline, StableDiffusionMixin): device, dtype, keep_model_device, + scheduler: KarrasDiffusionSchedulers, + unet: object, + timesteps, + scale_model_input_method, prompt_embeds: torch.Tensor, pooled_prompt_embeds: torch.Tensor, negative_prompt_embeds: torch.Tensor, @@ -158,6 +151,11 @@ class StableDiffusionXLFillPipeline(DiffusionPipeline, StableDiffusionMixin): ): self.controlnet = controlnet_model self._guidance_scale = guidance_scale + self.unet = unet + self.scheduler = scheduler + self.timesteps = timesteps + self.scale_model_input_method=scale_model_input_method + # 2. Define call parameters batch_size = 1 @@ -228,7 +226,7 @@ class StableDiffusionXLFillPipeline(DiffusionPipeline, StableDiffusionMixin): num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order ComfyUI_ProgressBar = ProgressBar(int(num_inference_steps)) - with self.progress_bar(total=num_inference_steps) as progress_bar: + with tqdm(total=num_inference_steps) as pbar: for i, t in enumerate(timesteps): # expand the latents if we are doing classifier free guidance latent_model_input = ( @@ -318,10 +316,5 @@ class StableDiffusionXLFillPipeline(DiffusionPipeline, StableDiffusionMixin): ComfyUI_ProgressBar.update(1) #yield latents_to_rgb(latents) - del self.unet - del self.controlnet - gc.collect() - torch.cuda.empty_cache() - latents = latents / 0.13025 yield latents