"""Image refiner: refiner pipeline, refiner container, and utilities. Contains: - RefinerImageProcessor: Image pre/post-processing for the diffusion pipeline - RefinerPipeline: DiffusionPipeline for image refinement - ImageRefinerContainer: nn.Module container for refiner sub-modules - IdentityWithArgs: Placeholder module for cond_proj - de_transform / tensor2pil: Tensor-to-PIL conversion utilities """ import inspect import math import warnings from dataclasses import dataclass from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np from safetensors.torch import load_file import torch import torch.nn as nn import torch.nn.functional as F from PIL import Image from diffusers import DiffusionPipeline from diffusers.configuration_utils import register_to_config from diffusers.image_processor import PipelineImageInput, VaeImageProcessor, is_valid_image_imagelist from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL from .refiner_modules import FlowMatchEulerDiscreteScheduler from .refiner_modules import Transformer2DModel, RotaryPosEmbed # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def _clean_config_dict(cfg, cls=None) -> dict: """Convert a PretrainedConfig to a clean dict for model construction. If ``cls`` is provided, only keeps keys that match the cls.__init__ params (allowlist approach). Otherwise falls back to blocklist filtering. """ if hasattr(cfg, "to_dict"): d = cfg.to_dict() elif isinstance(cfg, dict): d = dict(cfg) else: d = {k: v for k, v in vars(cfg).items()} if cls is not None: import inspect sig = inspect.signature(cls.__init__) valid_keys = set(sig.parameters.keys()) - {"self"} if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()): # Has **kwargs — can't filter by allowlist, fall through to blocklist pass else: return {k: v for k, v in d.items() if k in valid_keys} # Blocklist: remove HuggingFace PretrainedConfig metadata _PRETRAINED_CONFIG_KEYS = { "_name_or_path", "transformers_version", "model_type", "_commit_hash", "_attn_implementation", "_attn_implementation_autoset", "return_dict", "output_hidden_states", "output_attentions", "use_bfloat16", "torchscript", "torch_dtype", "is_encoder_decoder", "is_decoder", "add_cross_attention", "tie_encoder_decoder", "tie_word_embeddings", "cross_attention_hidden_size", "chunk_size_feed_forward", "decoder_start_token_id", "architectures", "finetuning_task", "id2label", "label2id", "prefix", "problem_type", "tokenizer_class", "task_specific_params", "pruned_heads", "bos_token_id", "eos_token_id", "pad_token_id", "sep_token_id", "max_length", "min_length", "do_sample", "early_stopping", "num_beams", "num_beam_groups", "diversity_penalty", "temperature", "top_k", "top_p", "typical_p", "repetition_penalty", "length_penalty", "no_repeat_ngram_size", "encoder_no_repeat_ngram_size", "bad_words_ids", "num_return_sequences", "output_scores", "return_dict_in_generate", "forced_bos_token_id", "forced_eos_token_id", "remove_invalid_values", "exponential_decay_length_penalty", "suppress_tokens", "begin_suppress_tokens", "tf_legacy_loss", "dtype", } return {k: v for k, v in d.items() if not k.startswith("_") and k not in _PRETRAINED_CONFIG_KEYS} # --------------------------------------------------------------------------- # Image Refiner Container (nn.Module for state_dict loading) # --------------------------------------------------------------------------- class ImageRefinerContainer(nn.Module): """Container for refiner components. Holds base_transformer, vae, cond_proj as nn.Module children so their parameters appear in the parent model's state_dict and are loaded automatically via from_pretrained. """ def __init__(self, visual_decoder_config): super().__init__() tc = visual_decoder_config.transformer_config vc = visual_decoder_config.vae_config self.base_transformer = Transformer2DModel(**_clean_config_dict(tc)) self.vae = AutoencoderKL(**_clean_config_dict(vc)) self.vae.requires_grad_(False) text_feat_dim = getattr(tc, "text_feat_dim", 3584) codebook_dim = getattr(visual_decoder_config, "codebook_dim", text_feat_dim) if codebook_dim != text_feat_dim: self.cond_proj = nn.Linear(codebook_dim, text_feat_dim) else: self.cond_proj = IdentityWithArgs() @classmethod def from_pretrained(cls, config, model_path: str): model = cls(config) weight_dict = load_file(model_path, device="cpu") model.load_state_dict({k.removeprefix("image_refiner."): v for k, v in weight_dict.items() if k.startswith("image_refiner.")}, strict=True) model.eval() return model @property def device(self): return next(self.parameters()).device @property def dtype(self): return next(self.parameters()).dtype class RefinerImageProcessor(VaeImageProcessor): """Image processor for refiner - extends diffusers' VaeImageProcessor.""" @register_to_config def __init__( self, do_resize: bool = True, vae_scale_factor: int = 16, resample: str = "lanczos", max_pixels: Optional[int] = None, max_side_length: Optional[int] = None, do_normalize: bool = True, do_binarize: bool = False, do_convert_grayscale: bool = False, ): super().__init__( do_resize=do_resize, vae_scale_factor=vae_scale_factor, resample=resample, do_normalize=do_normalize, do_binarize=do_binarize, do_convert_grayscale=do_convert_grayscale, ) self.max_pixels = max_pixels self.max_side_length = max_side_length def get_new_height_width( self, image: Union["PIL.Image.Image", np.ndarray, torch.Tensor], height: Optional[int] = None, width: Optional[int] = None, max_pixels: Optional[int] = None, max_side_length: Optional[int] = None, ) -> Tuple[int, int]: import PIL.Image if height is None: if isinstance(image, PIL.Image.Image): height = image.height elif isinstance(image, torch.Tensor): height = image.shape[2] else: height = image.shape[1] if width is None: if isinstance(image, PIL.Image.Image): width = image.width elif isinstance(image, torch.Tensor): width = image.shape[3] else: width = image.shape[2] if max_side_length is None: max_side_length = self.max_side_length if max_pixels is None: max_pixels = self.max_pixels ratio = 1.0 if max_side_length is not None: max_side_length_ratio = max_side_length / max(height, width) else: max_side_length_ratio = 1.0 cur_pixels = height * width max_pixels_ratio = (max_pixels / cur_pixels) ** 0.5 if max_pixels is not None else 1.0 ratio = min(max_pixels_ratio, max_side_length_ratio, 1.0) sf = self.config.vae_scale_factor new_height = int(height * ratio) // sf * sf new_width = int(width * ratio) // sf * sf return new_height, new_width def preprocess( self, image: PipelineImageInput, height: Optional[int] = None, width: Optional[int] = None, max_pixels: Optional[int] = None, max_side_length: Optional[int] = None, resize_mode: str = "default", crops_coords: Optional[Tuple[int, int, int, int]] = None, ) -> torch.Tensor: import PIL.Image supported_formats = (PIL.Image.Image, np.ndarray, torch.Tensor) if self.config.do_convert_grayscale and isinstance(image, (torch.Tensor, np.ndarray)) and image.ndim == 3: if isinstance(image, torch.Tensor): image = image.unsqueeze(1) else: if image.shape[-1] == 1: image = np.expand_dims(image, axis=0) else: image = np.expand_dims(image, axis=-1) if isinstance(image, list) and isinstance(image[0], np.ndarray) and image[0].ndim == 4: warnings.warn( "Passing `image` as a list of 4d np.ndarray is deprecated. " "Please concatenate the list along the batch dimension and pass it as a single 4d np.ndarray", FutureWarning, ) image = np.concatenate(image, axis=0) if isinstance(image, list) and isinstance(image[0], torch.Tensor) and image[0].ndim == 4: warnings.warn( "Passing `image` as a list of 4d torch.Tensor is deprecated. " "Please concatenate the list along the batch dimension and pass it as a single 4d torch.Tensor", FutureWarning, ) image = torch.cat(image, axis=0) if not is_valid_image_imagelist(image): raise ValueError( f"Input is in incorrect format. Currently, we only support " f"{', '.join(str(x) for x in supported_formats)}" ) if not isinstance(image, list): image = [image] if isinstance(image[0], PIL.Image.Image): if crops_coords is not None: image = [i.crop(crops_coords) for i in image] if self.config.do_resize: height, width = self.get_new_height_width(image[0], height, width, max_pixels, max_side_length) image = [self.resize(i, height, width, resize_mode=resize_mode) for i in image] if self.config.do_convert_grayscale: image = [self.convert_to_grayscale(i) for i in image] image = self.pil_to_numpy(image) image = self.numpy_to_pt(image) elif isinstance(image[0], np.ndarray): image = np.concatenate(image, axis=0) if image[0].ndim == 4 else np.stack(image, axis=0) image = self.numpy_to_pt(image) height, width = self.get_new_height_width(image, height, width, max_pixels, max_side_length) if self.config.do_resize: image = self.resize(image, height, width) elif isinstance(image[0], torch.Tensor): image = torch.cat(image, axis=0) if image[0].ndim == 4 else torch.stack(image, axis=0) if self.config.do_convert_grayscale and image.ndim == 3: image = image.unsqueeze(1) channel = image.shape[1] if channel == self.config.vae_latent_channels: return image height, width = self.get_new_height_width(image, height, width, max_pixels, max_side_length) if self.config.do_resize: image = self.resize(image, height, width) do_normalize = self.config.do_normalize if do_normalize and image.min() < 0: warnings.warn( "Passing `image` as torch tensor with value range in [-1,1] is deprecated. " f"The expected value range for image tensor is [0,1] when passing as pytorch tensor or numpy Array. " f"You passed `image` with value range [{image.min()},{image.max()}]", FutureWarning, ) do_normalize = False if do_normalize: image = self.normalize(image) if self.config.do_binarize: image = self.binarize(image) return image @dataclass class RefinerOutput: images: Union[List[Image.Image], torch.Tensor] class IdentityWithArgs(nn.Module): """Placeholder Identity module for cond_proj.""" def __init__(self, dtype=torch.float32, device=None): super().__init__() self.register_buffer("_dummy", torch.zeros((), dtype=dtype, device=device)) @property def dtype(self): return self._dummy.dtype @property def device(self): return self._dummy.device def forward(self, x, *args, **kwargs): return x def _retrieve_timesteps( scheduler: FlowMatchEulerDiscreteScheduler, num_inference_steps: Optional[int] = None, device: Optional[Union[str, torch.device]] = None, timesteps: Optional[List[int]] = None, **kwargs, ): # If scheduler uses dynamic shifting and caller passed num_tokens, compute mu # (same as training code refiner pipeline) num_tokens = kwargs.pop("num_tokens", None) if num_tokens is not None and getattr(scheduler.config, "use_dynamic_shifting", False): # Compute mu from num_tokens using scheduler's linear interpolation base_shift = getattr(scheduler.config, "base_shift", 0.5) max_shift = getattr(scheduler.config, "max_shift", 1.15) base_seq_len = getattr(scheduler.config, "base_image_seq_len", 256) max_seq_len = getattr(scheduler.config, "max_image_seq_len", 4096) m = (max_shift - base_shift) / (max_seq_len - base_seq_len) b = base_shift - m * base_seq_len mu = num_tokens * m + b kwargs["mu"] = mu accepted = set(inspect.signature(scheduler.set_timesteps).parameters.keys()) filtered_kwargs = {k: v for k, v in kwargs.items() if k in accepted} if timesteps is not None: if "timesteps" not in accepted: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" timestep schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps(timesteps=timesteps, device=device, **filtered_kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) else: scheduler.set_timesteps(num_inference_steps, device=device, **filtered_kwargs) timesteps = scheduler.timesteps return timesteps, num_inference_steps class RefinerPipeline(DiffusionPipeline): """ Image refiner evaluation pipeline. - cond comes from upstream model: encoder_hidden_states (quants / last_latent) - grid_thw_list is used to split cond (consistent with training) - image as ref image - Supports FlowMatchEulerDiscreteScheduler + velocity model """ def __init__( self, vae: AutoencoderKL, transformer: Transformer2DModel, scheduler: FlowMatchEulerDiscreteScheduler, cond_proj: Optional[nn.Module] = None, ): super().__init__() self.register_modules( vae=vae, transformer=transformer, scheduler=scheduler, cond_proj=cond_proj if cond_proj is not None else IdentityWithArgs(), ) self.vae_scale_factor = ( 2 ** (len(self.vae.config.block_out_channels) - 1) if hasattr(self.vae.config, "block_out_channels") else 8 ) self.image_processor = RefinerImageProcessor( vae_scale_factor=self.vae_scale_factor * 2, do_resize=True ) self.patch_size = int(getattr(self.transformer.config, "patch_size", 16)) self._num_timesteps: int = 0 self._current_timestep: Optional[torch.Tensor] = None self._interrupt: bool = False self._freqs_cis: Optional[torch.Tensor] = None self._text_guidance_scale: float = 1.0 self._image_guidance_scale: float = 1.0 self._cfg_range: Tuple[float, float] = (0.0, 1.0) @torch.no_grad() def _get_freqs_cis(self, device, dtype): if self._freqs_cis is None: self._freqs_cis = RotaryPosEmbed.get_freqs_cis( self.transformer.config.axes_dim_rope, self.transformer.config.axes_lens, theta=10000, ) return self._freqs_cis @staticmethod def _split_tokens( encoder_hidden_states: torch.Tensor, grid_thw_list: List[Tuple[int, int, int]], ) -> List[torch.Tensor]: splits = [int(h) * int(w) // 4 for (_, h, w) in grid_thw_list] return list(torch.split(encoder_hidden_states, splits, dim=1)) @staticmethod def _looks_like_latents(x: Union[torch.Tensor, Image.Image], latent_ch_hint: int = 16) -> bool: if not isinstance(x, torch.Tensor): return False if x.ndim not in (3, 4): return False c = int(x.shape[-3]) if c == 3: return False if c == latent_ch_hint: return True if c > 3 and c <= 32: return True return False @torch.no_grad() def _preprocess_to_vae_range(self, img: torch.Tensor) -> torch.Tensor: if img.dtype not in (torch.float32, torch.float16, torch.bfloat16): img = img.float() if img.max() > 1.5: img = img / 255.0 if img.min() >= 0.0 and img.max() <= 1.0: img = img * 2.0 - 1.0 return img.clamp(-1, 1) @torch.no_grad() def _encode_image_to_latents( self, img_any: Union[Image.Image, torch.Tensor], device, dtype, ) -> Tuple[torch.Tensor, int, int]: latent_ch_hint = int(getattr(getattr(self.vae, "config", None), "latent_channels", 16)) if self._looks_like_latents(img_any, latent_ch_hint=latent_ch_hint): z = img_any if z.ndim == 3: z = z.unsqueeze(0) z = z.to(device=device, dtype=dtype) H_lat, W_lat = z.shape[-2], z.shape[-1] return z, H_lat, W_lat if isinstance(img_any, Image.Image): img = torch.from_numpy( np.array(img_any).astype("float32") / 255.0 ).permute(2, 0, 1).unsqueeze(0) elif isinstance(img_any, torch.Tensor): img = img_any if img.ndim == 3: img = img.unsqueeze(0) else: raise TypeError("Unsupported image type. Use PIL.Image or torch.Tensor or latent Tensor.") img = self._preprocess_to_vae_range(img) H, W = img.shape[-2:] base = self.patch_size * self.vae_scale_factor target_H = max(base, math.ceil(H / base) * base) target_W = max(base, math.ceil(W / base) * base) if (H != target_H) or (W != target_W): img = F.interpolate(img, size=(target_H, target_W), mode="bilinear", align_corners=False) img = img.to(device=device, dtype=self.vae.dtype) posterior = self.vae.encode(img).latent_dist z0 = posterior.sample() if getattr(self.vae.config, "shift_factor", None) is not None: z0 = z0 - self.vae.config.shift_factor if getattr(self.vae.config, "scaling_factor", None) is not None: z0 = z0 * self.vae.config.scaling_factor z0 = z0.to(device=device, dtype=dtype) H_lat, W_lat = z0.shape[-2], z0.shape[-1] return z0, H_lat, W_lat @staticmethod def _expand_to_list(x, n): if x is None: return [None] * n if isinstance(x, (Image.Image, torch.Tensor)): return [x] * n assert isinstance(x, list), "`image` must be PIL / Tensor or list of them." assert len(x) == n, "`len(image)` must equal number of image chunks" return x @torch.no_grad() def _denoise_once( self, cond_tokens: torch.Tensor, ref_img: Optional[Union[Image.Image, torch.Tensor]], num_inference_steps: int = 28, timesteps: Optional[List[int]] = None, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, output_type: str = "pil", text_guidance_scale: float = 1.0, image_guidance_scale: float = 1.0, cfg_range: Tuple[float, float] = (0.0, 1.0), enable_processor_bar: bool = True, ): device = cond_tokens.device weight_dtype = self.transformer.dtype self._text_guidance_scale = text_guidance_scale self._image_guidance_scale = image_guidance_scale self._cfg_range = cfg_range cond_tokens = cond_tokens.to(device=device, dtype=weight_dtype) text_feats = self.cond_proj(cond_tokens) B, L, _ = text_feats.shape text_mask = torch.ones(B, L, device=device, dtype=torch.bool) ref_image_hidden_states = None H_lat: int W_lat: int if ref_img is not None: if isinstance(ref_img, torch.Tensor) and ref_img.ndim == 4 and ref_img.shape[0] == B: z_ref, H_lat, W_lat = self._encode_image_to_latents(ref_img, device=device, dtype=weight_dtype) ref_image_hidden_states = [[z_ref[b]] for b in range(B)] else: z_ref, H_lat, W_lat = self._encode_image_to_latents(ref_img, device=device, dtype=weight_dtype) z_single = z_ref[0] ref_image_hidden_states = [[z_single] for _ in range(B)] else: H_lat = W_lat = 128 // self.vae_scale_factor C_lat = getattr(self.transformer.config, "in_channels", None) if C_lat is None: if ref_image_hidden_states is not None: C_lat = ref_image_hidden_states[0][0].shape[0] else: raise ValueError("transformer.config.in_channels is None and no ref_img was provided.") latents_shape = (B, C_lat, H_lat, W_lat) if isinstance(generator, list): if len(generator) != B: raise ValueError( f"len(generator)={len(generator)} must equal B={B} when passing list of generators." ) latents = torch.stack( [ torch.randn( (1, C_lat, H_lat, W_lat), generator=generator[i], device=device, dtype=weight_dtype, ).squeeze(0) for i in range(B) ], dim=0, ) else: latents = torch.randn(latents_shape, generator=generator, device=device, dtype=weight_dtype) num_tokens = H_lat * W_lat timesteps_sched, num_inference_steps = _retrieve_timesteps( self.scheduler, num_inference_steps=num_inference_steps, device=device, timesteps=timesteps, num_tokens=num_tokens, ) num_warmup_steps = max(len(timesteps_sched) - num_inference_steps * self.scheduler.order, 0) self._num_timesteps = len(timesteps_sched) freqs_cis = self._get_freqs_cis(device=device, dtype=weight_dtype) progress_bar = self.progress_bar(total=num_inference_steps) if enable_processor_bar else None for i, t in enumerate(timesteps_sched): if self._interrupt: continue self._current_timestep = t timestep = t.expand(latents.shape[0]).to(latents.dtype) step_frac = i / max(len(timesteps_sched) - 1, 1) use_cfg = (cfg_range[0] <= step_frac <= cfg_range[1]) and ( text_guidance_scale > 1.0 or image_guidance_scale > 1.0 ) if not use_cfg: optional_kwargs: Dict[str, Any] = {} if "ref_image_hidden_states" in inspect.signature(self.transformer.forward).parameters: optional_kwargs["ref_image_hidden_states"] = ref_image_hidden_states model_pred = self.transformer( latents, timestep, text_feats, freqs_cis, text_mask, **optional_kwargs ) else: text_uncond = torch.zeros_like(text_feats) opt_kwargs_text: Dict[str, Any] = {} if "ref_image_hidden_states" in inspect.signature(self.transformer.forward).parameters: opt_kwargs_text["ref_image_hidden_states"] = ref_image_hidden_states model_pred_text = self.transformer( latents, timestep, text_feats, freqs_cis, text_mask, **opt_kwargs_text ) opt_kwargs_ref: Dict[str, Any] = {} if "ref_image_hidden_states" in inspect.signature(self.transformer.forward).parameters: opt_kwargs_ref["ref_image_hidden_states"] = ref_image_hidden_states model_pred_ref = self.transformer( latents, timestep, text_uncond, freqs_cis, text_mask, **opt_kwargs_ref ) opt_kwargs_uncond: Dict[str, Any] = {} if "ref_image_hidden_states" in inspect.signature(self.transformer.forward).parameters: opt_kwargs_uncond["ref_image_hidden_states"] = None model_pred_uncond = self.transformer( latents, timestep, text_uncond, freqs_cis, text_mask, **opt_kwargs_uncond ) if text_guidance_scale > 1.0 and image_guidance_scale > 1.0: model_pred = ( model_pred_uncond + image_guidance_scale * (model_pred_ref - model_pred_uncond) + text_guidance_scale * (model_pred_text - model_pred_ref) ) elif text_guidance_scale > 1.0: model_pred = model_pred_uncond + text_guidance_scale * (model_pred_text - model_pred_uncond) elif image_guidance_scale > 1.0: model_pred = model_pred_uncond + image_guidance_scale * (model_pred_ref - model_pred_uncond) else: model_pred = model_pred_text latents = self.scheduler.step(model_pred, t, latents, return_dict=False)[0] latents = latents.to(dtype=weight_dtype) if progress_bar is not None: if i == len(timesteps_sched) - 1 or ( (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0 ): progress_bar.update() if progress_bar is not None: progress_bar.close() self._current_timestep = None latents = latents.to(dtype=self.vae.dtype) if getattr(self.vae.config, "scaling_factor", None) is not None: latents = latents / self.vae.config.scaling_factor if getattr(self.vae.config, "shift_factor", None) is not None: latents = latents + self.vae.config.shift_factor image = self.vae.decode(latents, return_dict=False)[0] images = self.image_processor.postprocess(image, output_type=output_type) return images @torch.no_grad() def __call__( self, *, encoder_hidden_states: torch.Tensor, grid_thw_list: List[Tuple[int, int, int]], image: Union[Image.Image, torch.Tensor, List[Union[Image.Image, torch.Tensor]], None] = None, num_inference_steps: int = 28, timesteps: Optional[List[int]] = None, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, output_type: str = "pil", return_dict: bool = True, text_guidance_scale: float = 1.5, image_guidance_scale: float = 1.5, cfg_range: Tuple[float, float] = (0.0, 1.0), enable_processor_bar: bool = True, **kwargs, ) -> Union[RefinerOutput, List[Image.Image], torch.Tensor]: self._interrupt = False token_chunks = self._split_tokens(encoder_hidden_states, grid_thw_list) ref_list = self._expand_to_list(image, len(token_chunks)) results_pil: List[Image.Image] = [] results_pt: Optional[torch.Tensor] = None for tok, _, img_any in zip(token_chunks, grid_thw_list, ref_list): imgs = self._denoise_once( cond_tokens=tok, ref_img=img_any, num_inference_steps=num_inference_steps, timesteps=timesteps, generator=generator, output_type=output_type, text_guidance_scale=text_guidance_scale, image_guidance_scale=image_guidance_scale, cfg_range=cfg_range, enable_processor_bar=enable_processor_bar, ) if output_type == "pil": results_pil += imgs else: results_pt = imgs if results_pt is None else torch.cat([results_pt, imgs], dim=0) if not return_dict: return results_pil if output_type == "pil" else results_pt return RefinerOutput(images=results_pil if output_type == "pil" else results_pt) def de_transform( tensor: torch.Tensor, mean=(0.48145466, 0.4578275, 0.40821073), std=(0.26862954, 0.26130258, 0.27577711), rescale_factor: float = 1 / 255, ) -> torch.Tensor: """De-normalize and de-rescale, suitable for images processed by Qwen2VLImageProcessor.""" if tensor.ndim == 3: tensor = tensor.unsqueeze(0) mean_t = torch.tensor(mean).view(1, -1, 1, 1).to(tensor.device) std_t = torch.tensor(std).view(1, -1, 1, 1).to(tensor.device) tensor = tensor * std_t + mean_t tensor = tensor / rescale_factor tensor = torch.clamp(tensor / 255.0, 0, 1) return tensor def tensor2pil(image_t: torch.Tensor, image_mean, image_std) -> Image.Image: """Convert a tensor to a PIL Image.""" image_t = image_t.detach().cpu() rescale_factor = 1 / 255 sample = de_transform( image_t, mean=image_mean, std=image_std, rescale_factor=rescale_factor, )[0] ndarr = sample.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 0).to("cpu", torch.uint8).numpy() return Image.fromarray(ndarr)