| import inspect | |
| from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence | |
| import wandb | |
| from wandb.util import get_module | |
| if TYPE_CHECKING: | |
| np_array = get_module("numpy.array") | |
| torch_float_tensor = get_module("torch.FloatTensor") | |
| def chunkify(input_list, chunk_size) -> List: | |
| chunk_size = max(1, chunk_size) | |
| return [ | |
| input_list[i : i + chunk_size] for i in range(0, len(input_list), chunk_size) | |
| ] | |
| def get_updated_kwargs( | |
| pipeline: Any, args: Sequence[Any], kwargs: Dict[str, Any] | |
| ) -> Dict[str, Any]: | |
| pipeline_call_parameters = list( | |
| inspect.signature(pipeline.__call__).parameters.items() | |
| ) | |
| for idx, arg in enumerate(args): | |
| kwargs[pipeline_call_parameters[idx][0]] = arg | |
| for pipeline_parameter in pipeline_call_parameters: | |
| if pipeline_parameter[0] not in kwargs: | |
| kwargs[pipeline_parameter[0]] = pipeline_parameter[1].default | |
| if "generator" in kwargs: | |
| generator = kwargs["generator"] | |
| kwargs["generator"] = ( | |
| { | |
| "seed": generator.initial_seed(), | |
| "device": generator.device, | |
| "random_state": generator.get_state().cpu().numpy().tolist(), | |
| } | |
| if generator is not None | |
| else None | |
| ) | |
| if "ip_adapter_image" in kwargs: | |
| if kwargs["ip_adapter_image"] is not None: | |
| wandb.log({"IP-Adapter-Image": wandb.Image(kwargs["ip_adapter_image"])}) | |
| return kwargs | |
| def postprocess_pils_to_np(image: List) -> "np_array": | |
| np = get_module( | |
| "numpy", | |
| required="Please ensure NumPy is installed. You can run `pip install numpy` to install it.", | |
| ) | |
| return np.stack( | |
| [np.transpose(np.array(img).astype("uint8"), axes=(2, 0, 1)) for img in image], | |
| axis=0, | |
| ) | |
| def postprocess_np_arrays_for_video( | |
| images: List["np_array"], normalize: Optional[bool] = False | |
| ) -> "np_array": | |
| np = get_module( | |
| "numpy", | |
| required="Please ensure NumPy is installed. You can run `pip install numpy` to install it.", | |
| ) | |
| images = [(img * 255).astype("uint8") for img in images] if normalize else images | |
| return np.transpose(np.stack((images), axis=0), axes=(0, 3, 1, 2)) | |
| def decode_sdxl_t2i_latents(pipeline: Any, latents: "torch_float_tensor") -> List: | |
| """Decode latents generated by [`diffusers.StableDiffusionXLPipeline`](https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_xl#stable-diffusion-xl). | |
| Args: | |
| pipeline: (diffusers.DiffusionPipeline) The Diffusion Pipeline from | |
| [`diffusers`](https://huggingface.co/docs/diffusers). | |
| latents (torch.FloatTensor): The generated latents. | |
| Returns: | |
| List of `PIL` images corresponding to the generated latents. | |
| """ | |
| torch = get_module( | |
| "torch", | |
| required="Please ensure PyTorch is installed. You can check out https://pytorch.org/get-started/locally/#start-locally for installation instructions.", | |
| ) | |
| with torch.no_grad(): | |
| needs_upcasting = ( | |
| pipeline.vae.dtype == torch.float16 and pipeline.vae.config.force_upcast | |
| ) | |
| if needs_upcasting: | |
| pipeline.upcast_vae() | |
| latents = latents.to( | |
| next(iter(pipeline.vae.post_quant_conv.parameters())).dtype | |
| ) | |
| images = pipeline.vae.decode( | |
| latents / pipeline.vae.config.scaling_factor, return_dict=False | |
| )[0] | |
| if needs_upcasting: | |
| pipeline.vae.to(dtype=torch.float16) | |
| if pipeline.watermark is not None: | |
| images = pipeline.watermark.apply_watermark(images) | |
| images = pipeline.image_processor.postprocess(images, output_type="pil") | |
| pipeline.maybe_free_model_hooks() | |
| return images | |