| from typing import List, Optional, Union
|
|
|
| import torch
|
|
|
| from diffusers.image_processor import VaeImageProcessor
|
| from diffusers.pipelines.pipeline_utils import DiffusionPipeline, ImagePipelineOutput
|
| from diffusers.utils.torch_utils import randn_tensor
|
|
|
|
|
| class SiTPipeline(DiffusionPipeline):
|
| model_cpu_offload_seq = "transformer->vae"
|
|
|
| def __init__(self, transformer, scheduler, vae):
|
| super().__init__()
|
| self.register_modules(transformer=transformer, scheduler=scheduler, vae=vae)
|
| self.vae_scale_factor = 8
|
| self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
|
|
|
| @torch.no_grad()
|
| def __call__(
|
| self,
|
| class_labels: Union[int, List[int]] = 207,
|
| height: int = 256,
|
| width: int = 256,
|
| num_inference_steps: int = 250,
|
| guidance_scale: float = 4.0,
|
| generator: Optional[torch.Generator] = None,
|
| output_type: str = "pil",
|
| return_dict: bool = True,
|
| ):
|
| device = self._execution_device
|
| if isinstance(class_labels, int):
|
| class_labels = [class_labels]
|
| batch_size = len(class_labels)
|
|
|
| latent_h = height // self.vae_scale_factor
|
| latent_w = width // self.vae_scale_factor
|
| latents = randn_tensor(
|
| (batch_size, self.transformer.config.in_channels, latent_h, latent_w),
|
| generator=generator,
|
| device=device,
|
| dtype=self.transformer.dtype,
|
| )
|
|
|
| labels = torch.tensor(class_labels, device=device, dtype=torch.long)
|
| do_cfg = guidance_scale is not None and guidance_scale > 1.0
|
| if do_cfg:
|
| null_label = torch.full((batch_size,), self.transformer.config.num_classes, device=device, dtype=torch.long)
|
| labels = torch.cat([labels, null_label], dim=0)
|
|
|
| self.scheduler.set_timesteps(num_inference_steps, device=device)
|
| timesteps = self.scheduler.timesteps
|
|
|
| for t in self.progress_bar(timesteps):
|
| t_batch = torch.full((batch_size,), t, device=device, dtype=latents.dtype)
|
| model_input = latents
|
| if do_cfg:
|
| model_input = torch.cat([latents, latents], dim=0)
|
| t_batch = torch.cat([t_batch, t_batch], dim=0)
|
|
|
| model_pred = self.transformer(
|
| hidden_states=model_input,
|
| timestep=t_batch,
|
| class_labels=labels,
|
| ).sample
|
|
|
| if do_cfg:
|
| cond, uncond = model_pred.chunk(2, dim=0)
|
| model_pred = uncond + guidance_scale * (cond - uncond)
|
|
|
| latents = self.scheduler.step(model_pred, t, latents, generator=generator).prev_sample
|
|
|
| image = self.vae.decode(latents / 0.18215).sample
|
|
|
| if output_type == "pt":
|
| image = image
|
| else:
|
| image = self.image_processor.postprocess(image, output_type=output_type)
|
|
|
| if not return_dict:
|
| return (image,)
|
| return ImagePipelineOutput(images=image)
|
|
|
