|
| 1 | +from time import time |
| 2 | + |
| 3 | +import torch |
| 4 | +from diffusers import ( |
| 5 | + DiffusionPipeline, |
| 6 | + StableDiffusionXLImg2ImgPipeline, |
| 7 | +) |
| 8 | +from PIL import Image |
| 9 | + |
| 10 | +from backend.computing import Computing |
| 11 | +from backend.stablediffusion.modelmeta import ModelMeta |
| 12 | +from backend.stablediffusion.models.scheduler_types import SchedulerType |
| 13 | +from backend.stablediffusion.models.setting import ( |
| 14 | + StableDiffusionImageToImageSetting, |
| 15 | + StableDiffusionSetting, |
| 16 | +) |
| 17 | +from backend.stablediffusion.scheduler_mixin import SamplerMixin |
| 18 | + |
| 19 | + |
| 20 | +class StableDiffusionXl(SamplerMixin): |
| 21 | + def __init__(self, compute: Computing): |
| 22 | + self.compute = compute |
| 23 | + self.pipeline = None |
| 24 | + self.device = self.compute.name |
| 25 | + |
| 26 | + super().__init__() |
| 27 | + |
| 28 | + def get_text_to_image_xl_pipleline( |
| 29 | + self, |
| 30 | + model_id: str = "stabilityai/stable-diffusion-xl-base-1.0", |
| 31 | + low_vram_mode: bool = False, |
| 32 | + sampler: str = SchedulerType.DPMSolverMultistepScheduler.value, |
| 33 | + ): |
| 34 | + repo_id = model_id |
| 35 | + model_meta = ModelMeta(repo_id) |
| 36 | + is_lora_model = model_meta.is_loramodel() |
| 37 | + if is_lora_model: |
| 38 | + print("LoRA model detected") |
| 39 | + self.model_id = model_meta.get_lora_base_model() |
| 40 | + print(f"LoRA base model - {self.model_id}") |
| 41 | + else: |
| 42 | + self.model_id = model_id |
| 43 | + |
| 44 | + self.low_vram_mode = low_vram_mode |
| 45 | + print(f"StableDiffusion - {self.compute.name},{self.compute.datatype}") |
| 46 | + print(f"using model {model_id}") |
| 47 | + self.default_sampler = self.find_sampler( |
| 48 | + sampler, |
| 49 | + self.model_id, |
| 50 | + ) |
| 51 | + tic = time() |
| 52 | + self._load_model() |
| 53 | + delta = time() - tic |
| 54 | + print(f"Model loaded in {delta:.2f}s ") |
| 55 | + |
| 56 | + if self.pipeline is None: |
| 57 | + raise Exception("Text to image pipeline not initialized") |
| 58 | + if is_lora_model: |
| 59 | + self.pipeline.unet.load_attn_procs(repo_id) |
| 60 | + self._pipeline_to_device() |
| 61 | + components = self.pipeline.components |
| 62 | + self.img_to_img_pipeline = StableDiffusionXLImg2ImgPipeline(**components) |
| 63 | + |
| 64 | + def text_to_image_xl(self, setting: StableDiffusionSetting): |
| 65 | + if self.pipeline is None: |
| 66 | + raise Exception("Text to image pipeline not initialized") |
| 67 | + |
| 68 | + self.pipeline.scheduler = self.find_sampler( |
| 69 | + setting.scheduler, |
| 70 | + self.model_id, |
| 71 | + ) |
| 72 | + generator = None |
| 73 | + if setting.seed != -1: |
| 74 | + print(f"Using seed {setting.seed}") |
| 75 | + generator = torch.Generator(self.device).manual_seed(setting.seed) |
| 76 | + |
| 77 | + # if setting.attention_slicing: |
| 78 | + # self.pipeline.enable_attention_slicing() |
| 79 | + # else: |
| 80 | + # self.pipeline.disable_attention_slicing() |
| 81 | + |
| 82 | + if setting.vae_slicing: |
| 83 | + self.pipeline.enable_vae_slicing() |
| 84 | + else: |
| 85 | + self.pipeline.disable_vae_slicing() |
| 86 | + |
| 87 | + images = self.pipeline( |
| 88 | + setting.prompt, |
| 89 | + guidance_scale=setting.guidance_scale, |
| 90 | + num_inference_steps=setting.inference_steps, |
| 91 | + height=setting.image_height, |
| 92 | + width=setting.image_width, |
| 93 | + negative_prompt=setting.negative_prompt, |
| 94 | + num_images_per_prompt=setting.number_of_images, |
| 95 | + generator=generator, |
| 96 | + ).images |
| 97 | + |
| 98 | + # self.pipeline.unet = torch.compile( |
| 99 | + # self.pipeline.unet, |
| 100 | + # mode="reduce-overhead", |
| 101 | + # fullgraph=True, |
| 102 | + # ) |
| 103 | + return images |
| 104 | + |
| 105 | + def _pipeline_to_device(self): |
| 106 | + if self.low_vram_mode: |
| 107 | + print("Running in low VRAM mode,slower to generate images") |
| 108 | + self.pipeline.enable_sequential_cpu_offload() |
| 109 | + else: |
| 110 | + if self.compute.name == "cuda": |
| 111 | + self.pipeline = self.pipeline.to("cuda") |
| 112 | + elif self.compute.name == "mps": |
| 113 | + self.pipeline = self.pipeline.to("mps") |
| 114 | + |
| 115 | + def _load_full_precision_model(self): |
| 116 | + self.pipeline = DiffusionPipeline.from_pretrained( |
| 117 | + self.model_id, |
| 118 | + torch_dtype=self.compute.datatype, |
| 119 | + scheduler=self.default_sampler, |
| 120 | + ) |
| 121 | + |
| 122 | + def _load_model(self): |
| 123 | + if self.compute.name == "cuda": |
| 124 | + try: |
| 125 | + self.pipeline = DiffusionPipeline.from_pretrained( |
| 126 | + self.model_id, |
| 127 | + torch_dtype=self.compute.datatype, |
| 128 | + scheduler=self.default_sampler, |
| 129 | + use_safetensors=True, |
| 130 | + variant="fp16", |
| 131 | + ) |
| 132 | + except Exception as ex: |
| 133 | + print( |
| 134 | + f" The fp16 of the model not found using full precision model, {ex}" |
| 135 | + ) |
| 136 | + self._load_full_precision_model() |
| 137 | + else: |
| 138 | + self._load_full_precision_model() |
| 139 | + |
| 140 | + def image_to_image(self, setting: StableDiffusionImageToImageSetting): |
| 141 | + if setting.scheduler is None: |
| 142 | + raise Exception("Scheduler cannot be empty") |
| 143 | + |
| 144 | + print("Running image to image pipeline") |
| 145 | + self.img_to_img_pipeline.scheduler = self.find_sampler( # type: ignore |
| 146 | + setting.scheduler, |
| 147 | + self.model_id, |
| 148 | + ) |
| 149 | + generator = None |
| 150 | + if setting.seed != -1 and setting.seed: |
| 151 | + print(f"Using seed {setting.seed}") |
| 152 | + generator = torch.Generator(self.device).manual_seed(setting.seed) |
| 153 | + |
| 154 | + if setting.attention_slicing: |
| 155 | + self.img_to_img_pipeline.enable_attention_slicing() # type: ignore |
| 156 | + else: |
| 157 | + self.img_to_img_pipeline.disable_attention_slicing() # type: ignore |
| 158 | + |
| 159 | + if setting.vae_slicing: |
| 160 | + self.pipeline.enable_vae_slicing() # type: ignore |
| 161 | + else: |
| 162 | + self.pipeline.disable_vae_slicing() # type: ignore |
| 163 | + |
| 164 | + init_image = setting.image.resize( |
| 165 | + ( |
| 166 | + setting.image_width, |
| 167 | + setting.image_height, |
| 168 | + ), |
| 169 | + Image.Resampling.LANCZOS, |
| 170 | + ) |
| 171 | + images = self.img_to_img_pipeline( # type: ignore |
| 172 | + image=init_image, |
| 173 | + strength=setting.strength, |
| 174 | + prompt=setting.prompt, |
| 175 | + guidance_scale=setting.guidance_scale, |
| 176 | + num_inference_steps=setting.inference_steps, |
| 177 | + negative_prompt=setting.negative_prompt, |
| 178 | + num_images_per_prompt=setting.number_of_images, |
| 179 | + generator=generator, |
| 180 | + ).images |
| 181 | + return images |
0 commit comments