| from __future__ import annotations |
|
|
| import gc |
| import json |
| import tempfile |
| from typing import Generator |
|
|
| import numpy as np |
| import PIL.Image |
| import torch |
| from diffusers import DiffusionPipeline, StableDiffusionUpscalePipeline |
| from diffusers.pipelines.deepfloyd_if import (fast27_timesteps, |
| smart27_timesteps, |
| smart50_timesteps, |
| smart100_timesteps, |
| smart185_timesteps) |
|
|
| from settings import (DISABLE_AUTOMATIC_CPU_OFFLOAD, DISABLE_SD_X4_UPSCALER, |
| HF_TOKEN, MAX_NUM_IMAGES, MAX_NUM_STEPS, MAX_SEED, |
| RUN_GARBAGE_COLLECTION) |
|
|
|
|
| class Model: |
| def __init__(self): |
| self.device = torch.device( |
| 'cuda:0' if torch.cuda.is_available() else 'cpu') |
| self.pipe = None |
| self.super_res_1_pipe = None |
| self.super_res_2_pipe = None |
| self.watermark_image = None |
|
|
| if torch.cuda.is_available(): |
| self.load_weights() |
| self.watermark_image = PIL.Image.fromarray( |
| self.pipe.watermarker.watermark_image.to( |
| torch.uint8).cpu().numpy(), |
| mode='RGBA') |
|
|
| def load_weights(self) -> None: |
| self.pipe = DiffusionPipeline.from_pretrained( |
| 'DeepFloyd/IF-I-XL-v1.0', |
| torch_dtype=torch.float16, |
| variant='fp16', |
| use_safetensors=True, |
| use_auth_token=HF_TOKEN) |
| self.super_res_1_pipe = DiffusionPipeline.from_pretrained( |
| 'DeepFloyd/IF-II-L-v1.0', |
| text_encoder=None, |
| torch_dtype=torch.float16, |
| variant='fp16', |
| use_safetensors=True, |
| use_auth_token=HF_TOKEN) |
|
|
| if not DISABLE_SD_X4_UPSCALER: |
| self.super_res_2_pipe = StableDiffusionUpscalePipeline.from_pretrained( |
| 'stabilityai/stable-diffusion-x4-upscaler', |
| torch_dtype=torch.float16) |
|
|
| if DISABLE_AUTOMATIC_CPU_OFFLOAD: |
| self.pipe.to(self.device) |
| self.super_res_1_pipe.to(self.device) |
|
|
| self.pipe.unet.to(memory_format=torch.channels_last) |
| self.pipe.unet = torch.compile(self.pipe.unet, mode="reduce-overhead", fullgraph=True) |
|
|
| if not DISABLE_SD_X4_UPSCALER: |
| self.super_res_2_pipe.to(self.device) |
| else: |
| self.pipe.enable_model_cpu_offload() |
| self.super_res_1_pipe.enable_model_cpu_offload() |
| if not DISABLE_SD_X4_UPSCALER: |
| self.super_res_2_pipe.enable_model_cpu_offload() |
|
|
| def apply_watermark_to_sd_x4_upscaler_results( |
| self, images: list[PIL.Image.Image]) -> None: |
| w, h = images[0].size |
|
|
| stability_x4_upscaler_sample_size = 128 |
|
|
| coef = min(h / stability_x4_upscaler_sample_size, |
| w / stability_x4_upscaler_sample_size) |
| img_h, img_w = (int(h / coef), int(w / coef)) if coef < 1 else (h, w) |
|
|
| S1, S2 = 1024**2, img_w * img_h |
| K = (S2 / S1)**0.5 |
| watermark_size = int(K * 62) |
| watermark_x = img_w - int(14 * K) |
| watermark_y = img_h - int(14 * K) |
|
|
| watermark_image = self.watermark_image.copy().resize( |
| (watermark_size, watermark_size), |
| PIL.Image.Resampling.BICUBIC, |
| reducing_gap=None) |
|
|
| for image in images: |
| image.paste(watermark_image, |
| box=( |
| watermark_x - watermark_size, |
| watermark_y - watermark_size, |
| watermark_x, |
| watermark_y, |
| ), |
| mask=watermark_image.split()[-1]) |
|
|
| @staticmethod |
| def to_pil_images(images: torch.Tensor) -> list[PIL.Image.Image]: |
| images = (images / 2 + 0.5).clamp(0, 1) |
| images = images.cpu().permute(0, 2, 3, 1).float().numpy() |
| images = np.round(images * 255).astype(np.uint8) |
| return [PIL.Image.fromarray(image) for image in images] |
|
|
| @staticmethod |
| def check_seed(seed: int) -> None: |
| if not 0 <= seed <= MAX_SEED: |
| raise ValueError |
|
|
| @staticmethod |
| def check_num_images(num_images: int) -> None: |
| if not 1 <= num_images <= MAX_NUM_IMAGES: |
| raise ValueError |
|
|
| @staticmethod |
| def check_num_inference_steps(num_steps: int) -> None: |
| if not 1 <= num_steps <= MAX_NUM_STEPS: |
| raise ValueError |
|
|
| @staticmethod |
| def get_custom_timesteps(name: str) -> list[int] | None: |
| if name == 'none': |
| timesteps = None |
| elif name == 'fast27': |
| timesteps = fast27_timesteps |
| elif name == 'smart27': |
| timesteps = smart27_timesteps |
| elif name == 'smart50': |
| timesteps = smart50_timesteps |
| elif name == 'smart100': |
| timesteps = smart100_timesteps |
| elif name == 'smart185': |
| timesteps = smart185_timesteps |
| else: |
| raise ValueError |
| return timesteps |
|
|
| @staticmethod |
| def run_garbage_collection(): |
| gc.collect() |
| torch.cuda.empty_cache() |
|
|
| def run_stage1( |
| self, |
| prompt: str, |
| negative_prompt: str = '', |
| seed: int = 0, |
| num_images: int = 1, |
| guidance_scale_1: float = 7.0, |
| custom_timesteps_1: str = 'smart100', |
| num_inference_steps_1: int = 100, |
| ) -> tuple[list[PIL.Image.Image], str, str]: |
| self.check_seed(seed) |
| self.check_num_images(num_images) |
| self.check_num_inference_steps(num_inference_steps_1) |
|
|
| if RUN_GARBAGE_COLLECTION: |
| self.run_garbage_collection() |
|
|
| generator = torch.Generator(device=self.device).manual_seed(seed) |
|
|
| prompt_embeds, negative_embeds = self.pipe.encode_prompt( |
| prompt=prompt, negative_prompt=negative_prompt) |
|
|
| timesteps = self.get_custom_timesteps(custom_timesteps_1) |
|
|
| images = self.pipe(prompt_embeds=prompt_embeds, |
| negative_prompt_embeds=negative_embeds, |
| num_images_per_prompt=num_images, |
| guidance_scale=guidance_scale_1, |
| timesteps=timesteps, |
| num_inference_steps=num_inference_steps_1, |
| generator=generator, |
| output_type='pt').images |
| pil_images = self.to_pil_images(images) |
| self.pipe.watermarker.apply_watermark( |
| pil_images, self.pipe.unet.config.sample_size) |
|
|
| stage1_params = { |
| 'prompt': prompt, |
| 'negative_prompt': negative_prompt, |
| 'seed': seed, |
| 'num_images': num_images, |
| 'guidance_scale_1': guidance_scale_1, |
| 'custom_timesteps_1': custom_timesteps_1, |
| 'num_inference_steps_1': num_inference_steps_1, |
| } |
| with tempfile.NamedTemporaryFile(mode='w', delete=False) as param_file: |
| param_file.write(json.dumps(stage1_params)) |
| stage1_result = { |
| 'prompt_embeds': prompt_embeds, |
| 'negative_embeds': negative_embeds, |
| 'images': images, |
| 'pil_images': pil_images, |
| } |
| with tempfile.NamedTemporaryFile(delete=False) as result_file: |
| torch.save(stage1_result, result_file.name) |
| return pil_images, param_file.name, result_file.name |
|
|
| def run_stage2( |
| self, |
| stage1_result_path: str, |
| stage2_index: int, |
| seed_2: int = 0, |
| guidance_scale_2: float = 4.0, |
| custom_timesteps_2: str = 'smart50', |
| num_inference_steps_2: int = 50, |
| disable_watermark: bool = False, |
| ) -> PIL.Image.Image: |
| self.check_seed(seed_2) |
| self.check_num_inference_steps(num_inference_steps_2) |
|
|
| if RUN_GARBAGE_COLLECTION: |
| self.run_garbage_collection() |
|
|
| generator = torch.Generator(device=self.device).manual_seed(seed_2) |
|
|
| stage1_result = torch.load(stage1_result_path) |
| prompt_embeds = stage1_result['prompt_embeds'] |
| negative_embeds = stage1_result['negative_embeds'] |
| images = stage1_result['images'] |
| images = images[[stage2_index]] |
|
|
| timesteps = self.get_custom_timesteps(custom_timesteps_2) |
|
|
| out = self.super_res_1_pipe(image=images, |
| prompt_embeds=prompt_embeds, |
| negative_prompt_embeds=negative_embeds, |
| num_images_per_prompt=1, |
| guidance_scale=guidance_scale_2, |
| timesteps=timesteps, |
| num_inference_steps=num_inference_steps_2, |
| generator=generator, |
| output_type='pt', |
| noise_level=250).images |
| pil_images = self.to_pil_images(out) |
|
|
| if disable_watermark: |
| return pil_images[0] |
|
|
| self.super_res_1_pipe.watermarker.apply_watermark( |
| pil_images, self.super_res_1_pipe.unet.config.sample_size) |
| return pil_images[0] |
|
|
| def run_stage3( |
| self, |
| image: PIL.Image.Image, |
| prompt: str = '', |
| negative_prompt: str = '', |
| seed_3: int = 0, |
| guidance_scale_3: float = 9.0, |
| num_inference_steps_3: int = 75, |
| ) -> PIL.Image.Image: |
| self.check_seed(seed_3) |
| self.check_num_inference_steps(num_inference_steps_3) |
|
|
| if RUN_GARBAGE_COLLECTION: |
| self.run_garbage_collection() |
|
|
| generator = torch.Generator(device=self.device).manual_seed(seed_3) |
| out = self.super_res_2_pipe(image=image, |
| prompt=prompt, |
| negative_prompt=negative_prompt, |
| num_images_per_prompt=1, |
| guidance_scale=guidance_scale_3, |
| num_inference_steps=num_inference_steps_3, |
| generator=generator, |
| noise_level=100).images |
| self.apply_watermark_to_sd_x4_upscaler_results(out) |
| return out[0] |
|
|
| def run_stage2_3( |
| self, |
| stage1_result_path: str, |
| stage2_index: int, |
| seed_2: int = 0, |
| guidance_scale_2: float = 4.0, |
| custom_timesteps_2: str = 'smart50', |
| num_inference_steps_2: int = 50, |
| prompt: str = '', |
| negative_prompt: str = '', |
| seed_3: int = 0, |
| guidance_scale_3: float = 9.0, |
| num_inference_steps_3: int = 75, |
| ) -> Generator[PIL.Image.Image]: |
| self.check_seed(seed_3) |
| self.check_num_inference_steps(num_inference_steps_3) |
|
|
| out_image = self.run_stage2( |
| stage1_result_path=stage1_result_path, |
| stage2_index=stage2_index, |
| seed_2=seed_2, |
| guidance_scale_2=guidance_scale_2, |
| custom_timesteps_2=custom_timesteps_2, |
| num_inference_steps_2=num_inference_steps_2, |
| disable_watermark=True) |
| temp_image = out_image.copy() |
| self.super_res_1_pipe.watermarker.apply_watermark( |
| [temp_image], self.super_res_1_pipe.unet.config.sample_size) |
| yield temp_image |
| yield self.run_stage3(image=out_image, |
| prompt=prompt, |
| negative_prompt=negative_prompt, |
| seed_3=seed_3, |
| guidance_scale_3=guidance_scale_3, |
| num_inference_steps_3=num_inference_steps_3) |
|
|
