File size: 6,803 Bytes
2a4c86a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | from dataclasses import dataclass
from pathlib import Path
from typing import List, Tuple
import numpy as np
import torch
from diffusers import DiffusionPipeline
from diffusers.pipelines.pipeline_utils import ImagePipelineOutput
from diffusers.utils import BaseOutput
from diffusers.utils.torch_utils import randn_tensor
from .modeling_jit_transformer_2d import JiTTransformer2DModel
from .scheduling_jit import JiTScheduler
@dataclass
class JiTPipelineOutput(BaseOutput):
images: List["PIL.Image.Image"] | np.ndarray | torch.Tensor
class JiTPipeline(DiffusionPipeline):
model_cpu_offload_seq = "transformer"
def __init__(self, transformer: JiTTransformer2DModel, scheduler: JiTScheduler | None = None):
super().__init__()
self.register_modules(transformer=transformer, scheduler=scheduler or JiTScheduler())
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
model_kwargs = dict(kwargs)
transformer_subfolder = model_kwargs.pop("transformer_subfolder", None)
scheduler_subfolder = model_kwargs.pop("scheduler_subfolder", None)
scheduler_kwargs = model_kwargs.pop("scheduler_kwargs", {})
if transformer_subfolder is not None:
transformer_path = str(Path(pretrained_model_name_or_path) / transformer_subfolder)
else:
transformer_path = pretrained_model_name_or_path
transformer = JiTTransformer2DModel.from_pretrained(transformer_path, **model_kwargs)
try:
scheduler = JiTScheduler.from_pretrained(
pretrained_model_name_or_path,
subfolder=scheduler_subfolder,
**scheduler_kwargs,
)
except Exception:
scheduler = JiTScheduler(**scheduler_kwargs)
return cls(transformer=transformer, scheduler=scheduler)
@torch.no_grad()
def __call__(
self,
class_labels: int | List[int] | torch.Tensor,
num_inference_steps: int = 50,
guidance_scale: float = 2.9,
guidance_interval_min: float = 0.1,
guidance_interval_max: float = 1.0,
noise_scale: float = 2.0,
t_eps: float = 5e-2,
sampling_method: str | None = None,
generator: torch.Generator | List[torch.Generator] | None = None,
output_type: str = "pil",
return_dict: bool = True,
) -> JiTPipelineOutput | ImagePipelineOutput | Tuple:
if output_type not in {"pil", "np", "pt"}:
raise ValueError("output_type must be one of: 'pil', 'np', 'pt'.")
if sampling_method is not None and sampling_method not in {"heun", "euler"}:
raise ValueError("sampling_method must be one of: 'heun', 'euler'.")
if num_inference_steps < 2:
raise ValueError("num_inference_steps must be >= 2.")
if sampling_method is not None and sampling_method != self.scheduler.config.solver:
self.scheduler = JiTScheduler.from_config(self.scheduler.config, solver=sampling_method)
if isinstance(class_labels, int):
class_labels = [class_labels]
if isinstance(class_labels, list):
class_labels = torch.tensor(class_labels, device=self._execution_device, dtype=torch.long)
else:
class_labels = class_labels.to(self._execution_device, dtype=torch.long).reshape(-1)
batch_size = class_labels.shape[0]
latent_size = int(self.transformer.config.sample_size)
latent_channels = int(getattr(self.transformer.config, "in_channels", 3))
num_classes = int(self.transformer.config.num_class_embeds)
class_labels = class_labels.clamp(0, num_classes - 1)
class_null = torch.full_like(class_labels, num_classes)
latents = randn_tensor(
shape=(batch_size, latent_channels, latent_size, latent_size),
generator=generator,
device=self._execution_device,
dtype=self.transformer.dtype,
) * noise_scale
self.scheduler.set_timesteps(num_inference_steps=num_inference_steps, device=self._execution_device)
timesteps = self.scheduler.timesteps.to(device=self._execution_device, dtype=latents.dtype)
def forward_cfg(z_value: torch.Tensor, t: torch.Tensor | float) -> torch.Tensor:
t = torch.as_tensor(t, device=self._execution_device, dtype=latents.dtype)
x_cond = self.transformer(sample=z_value, timestep=t.flatten(), class_labels=class_labels).sample
v_cond = (x_cond - z_value) / (1.0 - t).clamp_min(t_eps)
x_uncond = self.transformer(sample=z_value, timestep=t.flatten(), class_labels=class_null).sample
v_uncond = (x_uncond - z_value) / (1.0 - t).clamp_min(t_eps)
interval_mask = (t < guidance_interval_max) & (t > guidance_interval_min)
scale = torch.where(
interval_mask,
torch.tensor(guidance_scale, device=self._execution_device, dtype=latents.dtype),
torch.tensor(1.0, device=self._execution_device, dtype=latents.dtype),
)
return v_uncond + scale * (v_cond - v_uncond)
for i in self.progress_bar(range(num_inference_steps - 1)):
t, t_next = timesteps[i], timesteps[i + 1]
model_output = forward_cfg(latents, t)
if self.scheduler.config.solver == "heun":
latents = self.scheduler.step(
model_output=model_output,
timestep=t,
next_timestep=t_next,
sample=latents,
model_fn=forward_cfg,
).prev_sample
else:
latents = self.scheduler.step(
model_output=model_output,
timestep=t,
next_timestep=t_next,
sample=latents,
).prev_sample
# Match the original JiT implementation: always use Euler for the final step.
t, t_next = timesteps[-2], timesteps[-1]
model_output = forward_cfg(latents, t)
latents = self.scheduler.euler_step(
model_output=model_output,
timestep=t,
next_timestep=t_next,
sample=latents,
).prev_sample
images_pt = ((latents.float().clamp(-1, 1) + 1.0) / 2.0).cpu()
if output_type == "pt":
images = images_pt
else:
images_np = images_pt.permute(0, 2, 3, 1).numpy()
if output_type == "np":
images = images_np
else:
images = self.numpy_to_pil(images_np)
self.maybe_free_model_hooks()
if not return_dict:
return (images,)
return JiTPipelineOutput(images=images)
|