# Copyright 2024 NVIDIA and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 import numpy as np from typing import Callable, Dict, List, Optional, Tuple, Union import torch from PIL import Image from diffusers.utils.torch_utils import randn_tensor try: from .pipeline_pixeldit import PixelDiTPipeline from .pipeline_output import PixelDiTPipelineOutput except ImportError: from pipeline_pixeldit import PixelDiTPipeline from pipeline_output import PixelDiTPipelineOutput def _to_pixel_tensor(image, width, height, device, dtype): """Convert a PIL Image or float tensor to [B, 3, H, W] in [-1, 1].""" if isinstance(image, Image.Image): image = image.convert("RGB").resize((width, height)) image = np.array(image, dtype=np.float32) image = torch.from_numpy(image).permute(2, 0, 1).div(127.5).sub(1.0) image = image.unsqueeze(0) elif isinstance(image, np.ndarray): if image.dtype == np.uint8: image = torch.from_numpy(image.astype(np.float32)).div(127.5).sub(1.0) else: image = torch.from_numpy(image.astype(np.float32)).mul(2.0).sub(1.0) if image.dim() == 3: image = image.permute(2, 0, 1).unsqueeze(0) elif isinstance(image, torch.Tensor): if image.dim() == 3: image = image.unsqueeze(0) if image.is_floating_point() and image.max() <= 1.0 + 1e-4: image = image.mul(2.0).sub(1.0) return image.to(device=device, dtype=dtype) class PixelDiTImg2ImgPipeline(PixelDiTPipeline): """ Img2img pipeline for PixelDiT. Inherits everything from :class:`PixelDiTPipeline` — same model, same text encoder, same LoRA API, same schedulers. Pass an input image and a ``strength`` value to control how much the image is modified: ``strength=1.0`` equals pure text-to-image generation; ``strength=0.1`` barely changes the input. Because PixelDiT is a pixel-space model (no VAE), noise is injected directly on the pixel tensor using the flow-matching formula: ``x_t = (1 − σ) · image + σ · noise`` Note: PixelDiT needs ≥ 45 total denoising steps for clean output. With low ``strength`` the effective step count drops — keep ``num_inference_steps`` at 50+ to compensate. Example:: from diffusers.pipelines.pixeldit import PixelDiTImg2ImgPipeline from PIL import Image import torch pipe = PixelDiTImg2ImgPipeline.from_pretrained( "madtune/pixeldit-diffusers", torch_dtype=torch.bfloat16 ) pipe.to("cuda") init = Image.open("photo.jpg").convert("RGB") out = pipe( prompt="a cinematic landscape, golden hour", image=init, strength=0.75, num_inference_steps=50, ).images[0] out.save("img2img_out.png") """ @classmethod def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): """ Load from the same HF repo as :class:`PixelDiTPipeline`. Internally loads a T2I pipeline, then transfers its components into this class. """ import diffusers from .modeling_pixeldit_hf import PixelDiTModel if not hasattr(diffusers, "PixelDiTModel"): diffusers.PixelDiTModel = PixelDiTModel t2i = PixelDiTPipeline.from_pretrained(pretrained_model_name_or_path, **kwargs) return cls( transformer=t2i.transformer, scheduler=t2i.scheduler, text_encoder=t2i.text_encoder, tokenizer=t2i.tokenizer, ) @torch.no_grad() def __call__( self, prompt: Union[str, List[str]], image: Union[Image.Image, torch.Tensor, np.ndarray], strength: float = 0.8, negative_prompt: Optional[Union[str, List[str]]] = None, height: int = 512, width: int = 512, num_inference_steps: int = 20, guidance_scale: float = 3.5, flow_shift: Optional[float] = None, generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None, output_type: str = "pil", return_dict: bool = True, cross_attention_kwargs: Optional[Dict[str, any]] = None, callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None, callback_on_step_end_tensor_inputs: List[str] = ["latents"], **kwargs, ) -> Union[PixelDiTPipelineOutput, Tuple]: """ Args: prompt: Text prompt(s) guiding image generation. image: Input image. Accepts PIL ``Image``, ``numpy.ndarray`` (H×W×3 uint8 or float), or ``torch.Tensor`` (3×H×W or B×3×H×W). strength: How much to transform the input (0 < strength ≤ 1). ``1.0`` = full noise (equivalent to t2i). Recommended range: 0.5–0.85. negative_prompt: Optional negative prompt(s). height: Output height in pixels (must be divisible by 16). width: Output width in pixels (must be divisible by 16). num_inference_steps: Total scheduler steps. Use ≥ 50 for best quality. guidance_scale: CFG scale. ~3.5–7.5 works well. flow_shift: Override the scheduler's flow shift at runtime (e.g. 3.0 for 512px, 4.0 for 1024px). Leaves the scheduler config unchanged if ``None``. generator: Torch RNG for reproducibility. output_type: ``"pil"`` (default) or ``"np"`` (uint8 numpy array). return_dict: If ``True`` returns :class:`PixelDiTPipelineOutput`, else a tuple. cross_attention_kwargs: Passed to the attention processor (e.g. ``{"scale": 0.8}`` to adjust LoRA strength at inference). callback_on_step_end: Optional callable invoked at the end of each denoising step. callback_on_step_end_tensor_inputs: Names of tensors forwarded to the callback. Returns: :class:`PixelDiTPipelineOutput` or ``tuple``. """ device = self._execution_device dtype = self.transformer.dtype self._guidance_scale = guidance_scale lora_scale = (cross_attention_kwargs or {}).get("scale", None) if isinstance(prompt, str): prompt = [prompt] batch_size = len(prompt) self.check_inputs(prompt, height, width, negative_prompt) prompt_embeds, negative_prompt_embeds = self.encode_prompt( prompt, device=device, dtype=dtype, do_classifier_free_guidance=self.do_classifier_free_guidance, negative_prompt=negative_prompt, lora_scale=lora_scale, ) # Override flow shift if requested (reverts after this call via set_timesteps) if flow_shift is not None: self.scheduler.config.shift = flow_shift self.scheduler.set_timesteps(num_inference_steps, device=device) timesteps = self.scheduler.timesteps self._num_timesteps = len(timesteps) # Skip to the start timestep determined by strength t_start = max(0, int(num_inference_steps * (1.0 - strength))) timesteps = timesteps[t_start:] if len(timesteps) == 0: raise ValueError( f"strength={strength} with num_inference_steps={num_inference_steps} " "produces 0 denoising steps. Increase strength or num_inference_steps." ) # Preprocess image and add flow-matching noise at sigma_start img_tensor = _to_pixel_tensor(image, width, height, device, dtype) if img_tensor.shape[0] == 1 and batch_size > 1: img_tensor = img_tensor.expand(batch_size, -1, -1, -1).contiguous() sigma_start = timesteps[0].float() / 1000.0 noise = randn_tensor(img_tensor.shape, generator=generator, device=device, dtype=dtype) latents = (1.0 - sigma_start) * img_tensor + sigma_start * noise # Denoising loop for i, t in enumerate(self.progress_bar(timesteps)): if self.do_classifier_free_guidance: latent_model_input = torch.cat([latents] * 2) embeds = torch.cat([negative_prompt_embeds, prompt_embeds]) else: latent_model_input = latents embeds = prompt_embeds t_input = t.expand(latent_model_input.shape[0]) noise_pred = self.transformer(latent_model_input, t_input, embeds) if self.do_classifier_free_guidance: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) if hasattr(self.scheduler, "scale_model_input"): latents = self.scheduler.step( noise_pred, t, self.scheduler.scale_model_input(latents, t), return_dict=False, )[0] else: latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] if callback_on_step_end is not None: cb_kwargs = {k: locals()[k] for k in callback_on_step_end_tensor_inputs} callback_on_step_end(i, t, cb_kwargs) # Decode (pixel-space — just clamp and normalise) image_out = (latents.clamp(-1, 1) + 1) / 2 image_out = (image_out * 255).byte().permute(0, 2, 3, 1).cpu().numpy() if output_type == "pil": image_out = [Image.fromarray(img) for img in image_out] self.maybe_free_model_hooks() if not return_dict: return (image_out,) return PixelDiTPipelineOutput(images=image_out)