Allex21/LT
0
1import logging2import sys3import threading4import torch5from torchvision import transforms6from typing import *7from diffusers import EulerAncestralDiscreteScheduler8import diffusers.schedulers.scheduling_euler_ancestral_discrete9from diffusers.schedulers.scheduling_euler_ancestral_discrete import EulerAncestralDiscreteSchedulerOutput10import cv211from PIL import Image12import numpy as np13 14 15def fire_in_thread(f, *args, **kwargs):16 threading.Thread(target=f, args=args, kwargs=kwargs).start()17 18 19def add_logging_arguments(parser):20 parser.add_argument(21 "--console_log_level",22 type=str,23 default=None,24 choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],25 help="Set the logging level, default is INFO / ログレベルを設定する。デフォルトはINFO",26 )27 parser.add_argument(28 "--console_log_file",29 type=str,30 default=None,31 help="Log to a file instead of stderr / 標準エラー出力ではなくファイルにログを出力する",32 )33 parser.add_argument("--console_log_simple", action="store_true", help="Simple log output / シンプルなログ出力")34 35 36def setup_logging(args=None, log_level=None, reset=False):37 if logging.root.handlers:38 if reset:39 # remove all handlers40 for handler in logging.root.handlers[:]:41 logging.root.removeHandler(handler)42 else:43 return44 45 # log_level can be set by the caller or by the args, the caller has priority. If not set, use INFO46 if log_level is None and args is not None:47 log_level = args.console_log_level48 if log_level is None:49 log_level = "INFO"50 log_level = getattr(logging, log_level)51 52 msg_init = None53 if args is not None and args.console_log_file:54 handler = logging.FileHandler(args.console_log_file, mode="w")55 else:56 handler = None57 if not args or not args.console_log_simple:58 try:59 from rich.logging import RichHandler60 from rich.console import Console61 from rich.logging import RichHandler62 63 handler = RichHandler(console=Console(stderr=True))64 except ImportError:65 # print("rich is not installed, using basic logging")66 msg_init = "rich is not installed, using basic logging"67 68 if handler is None:69 handler = logging.StreamHandler(sys.stdout) # same as print70 handler.propagate = False71 72 formatter = logging.Formatter(73 fmt="%(message)s",74 datefmt="%Y-%m-%d %H:%M:%S",75 )76 handler.setFormatter(formatter)77 logging.root.setLevel(log_level)78 logging.root.addHandler(handler)79 80 if msg_init is not None:81 logger = logging.getLogger(__name__)82 logger.info(msg_init)83 84 85def pil_resize(image, size, interpolation=Image.LANCZOS):86 has_alpha = image.shape[2] == 4 if len(image.shape) == 3 else False87 88 if has_alpha:89 pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGRA2RGBA))90 else:91 pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))92 93 resized_pil = pil_image.resize(size, interpolation)94 95 # Convert back to cv2 format96 if has_alpha:97 resized_cv2 = cv2.cvtColor(np.array(resized_pil), cv2.COLOR_RGBA2BGRA)98 else:99 resized_cv2 = cv2.cvtColor(np.array(resized_pil), cv2.COLOR_RGB2BGR)100 101 return resized_cv2102 103 104# TODO make inf_utils.py105 106 107# region Gradual Latent hires fix108 109 110class GradualLatent:111 def __init__(112 self,113 ratio,114 start_timesteps,115 every_n_steps,116 ratio_step,117 s_noise=1.0,118 gaussian_blur_ksize=None,119 gaussian_blur_sigma=0.5,120 gaussian_blur_strength=0.5,121 unsharp_target_x=True,122 ):123 self.ratio = ratio124 self.start_timesteps = start_timesteps125 self.every_n_steps = every_n_steps126 self.ratio_step = ratio_step127 self.s_noise = s_noise128 self.gaussian_blur_ksize = gaussian_blur_ksize129 self.gaussian_blur_sigma = gaussian_blur_sigma130 self.gaussian_blur_strength = gaussian_blur_strength131 self.unsharp_target_x = unsharp_target_x132 133 def __str__(self) -> str:134 return (135 f"GradualLatent(ratio={self.ratio}, start_timesteps={self.start_timesteps}, "136 + f"every_n_steps={self.every_n_steps}, ratio_step={self.ratio_step}, s_noise={self.s_noise}, "137 + f"gaussian_blur_ksize={self.gaussian_blur_ksize}, gaussian_blur_sigma={self.gaussian_blur_sigma}, gaussian_blur_strength={self.gaussian_blur_strength}, "138 + f"unsharp_target_x={self.unsharp_target_x})"139 )140 141 def apply_unshark_mask(self, x: torch.Tensor):142 if self.gaussian_blur_ksize is None:143 return x144 blurred = transforms.functional.gaussian_blur(x, self.gaussian_blur_ksize, self.gaussian_blur_sigma)145 # mask = torch.sigmoid((x - blurred) * self.gaussian_blur_strength)146 mask = (x - blurred) * self.gaussian_blur_strength147 sharpened = x + mask148 return sharpened149 150 def interpolate(self, x: torch.Tensor, resized_size, unsharp=True):151 org_dtype = x.dtype152 if org_dtype == torch.bfloat16:153 x = x.float()154 155 x = torch.nn.functional.interpolate(x, size=resized_size, mode="bicubic", align_corners=False).to(dtype=org_dtype)156 157 # apply unsharp mask / アンシャープマスクを適用する158 if unsharp and self.gaussian_blur_ksize:159 x = self.apply_unshark_mask(x)160 161 return x162 163 164class EulerAncestralDiscreteSchedulerGL(EulerAncestralDiscreteScheduler):165 def __init__(self, *args, **kwargs):166 super().__init__(*args, **kwargs)167 self.resized_size = None168 self.gradual_latent = None169 170 def set_gradual_latent_params(self, size, gradual_latent: GradualLatent):171 self.resized_size = size172 self.gradual_latent = gradual_latent173 174 def step(175 self,176 model_output: torch.FloatTensor,177 timestep: Union[float, torch.FloatTensor],178 sample: torch.FloatTensor,179 generator: Optional[torch.Generator] = None,180 return_dict: bool = True,181 ) -> Union[EulerAncestralDiscreteSchedulerOutput, Tuple]:182 """183 Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion184 process from the learned model outputs (most often the predicted noise).185 186 Args:187 model_output (`torch.FloatTensor`):188 The direct output from learned diffusion model.189 timestep (`float`):190 The current discrete timestep in the diffusion chain.191 sample (`torch.FloatTensor`):192 A current instance of a sample created by the diffusion process.193 generator (`torch.Generator`, *optional*):194 A random number generator.195 return_dict (`bool`):196 Whether or not to return a197 [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] or tuple.198 199 Returns:200 [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] or `tuple`:201 If return_dict is `True`,202 [`~schedulers.scheduling_euler_ancestral_discrete.EulerAncestralDiscreteSchedulerOutput`] is returned,203 otherwise a tuple is returned where the first element is the sample tensor.204 205 """206 207 if isinstance(timestep, int) or isinstance(timestep, torch.IntTensor) or isinstance(timestep, torch.LongTensor):208 raise ValueError(209 (210 "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to"211 " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass"212 " one of the `scheduler.timesteps` as a timestep."213 ),214 )215 216 if not self.is_scale_input_called:217 # logger.warning(218 print(219 "The `scale_model_input` function should be called before `step` to ensure correct denoising. "220 "See `StableDiffusionPipeline` for a usage example."221 )222 223 if self.step_index is None:224 self._init_step_index(timestep)225 226 sigma = self.sigmas[self.step_index]227 228 # 1. compute predicted original sample (x_0) from sigma-scaled predicted noise229 if self.config.prediction_type == "epsilon":230 pred_original_sample = sample - sigma * model_output231 elif self.config.prediction_type == "v_prediction":232 # * c_out + input * c_skip233 pred_original_sample = model_output * (-sigma / (sigma**2 + 1) ** 0.5) + (sample / (sigma**2 + 1))234 elif self.config.prediction_type == "sample":235 raise NotImplementedError("prediction_type not implemented yet: sample")236 else:237 raise ValueError(f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, or `v_prediction`")238 239 sigma_from = self.sigmas[self.step_index]240 sigma_to = self.sigmas[self.step_index + 1]241 sigma_up = (sigma_to**2 * (sigma_from**2 - sigma_to**2) / sigma_from**2) ** 0.5242 sigma_down = (sigma_to**2 - sigma_up**2) ** 0.5243 244 # 2. Convert to an ODE derivative245 derivative = (sample - pred_original_sample) / sigma246 247 dt = sigma_down - sigma248 249 device = model_output.device250 if self.resized_size is None:251 prev_sample = sample + derivative * dt252 253 noise = diffusers.schedulers.scheduling_euler_ancestral_discrete.randn_tensor(254 model_output.shape, dtype=model_output.dtype, device=device, generator=generator255 )256 s_noise = 1.0257 else:258 print("resized_size", self.resized_size, "model_output.shape", model_output.shape, "sample.shape", sample.shape)259 s_noise = self.gradual_latent.s_noise260 261 if self.gradual_latent.unsharp_target_x:262 prev_sample = sample + derivative * dt263 prev_sample = self.gradual_latent.interpolate(prev_sample, self.resized_size)264 else:265 sample = self.gradual_latent.interpolate(sample, self.resized_size)266 derivative = self.gradual_latent.interpolate(derivative, self.resized_size, unsharp=False)267 prev_sample = sample + derivative * dt268 269 noise = diffusers.schedulers.scheduling_euler_ancestral_discrete.randn_tensor(270 (model_output.shape[0], model_output.shape[1], self.resized_size[0], self.resized_size[1]),271 dtype=model_output.dtype,272 device=device,273 generator=generator,274 )275 276 prev_sample = prev_sample + noise * sigma_up * s_noise277 278 # upon completion increase step index by one279 self._step_index += 1280 281 if not return_dict:282 return (prev_sample,)283 284 return EulerAncestralDiscreteSchedulerOutput(prev_sample=prev_sample, pred_original_sample=pred_original_sample)285 286 287# endregion288