_diffusion.py

# Copyright 2025 Google LLC
#
# 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
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import functools

import torch
from time import time
from diffusers import DiffusionPipeline
from torch.utils import _pytree as pytree

import torchax
import torchax.functions
from torchax.extra import torch_view, jax_view

import jax
import torch.func


class CompiledModule:
  def __init__(self, model):
    weights = model.state_dict()
    weights.update(model.named_parameters())
    self._weights = pytree.tree_map_only(
      torch.Tensor, torchax.tensor.move_to_device, weights
    )
    self._model = model

    self._func_jitted_torch = None  # torch_view(func_mod_jitted)

  def _maybe_move_tensor(self, tensor):
    if isinstance(tensor, torch.Tensor) and not isinstance(
      tensor, torchax.tensor.Tensor
    ):
      return torchax.tensor.move_to_device(tensor)
    return tensor

  def _make_jitted(self, args, kwargs):
    static = []
    for i, a in enumerate(args):
      if not isinstance(a, torch.Tensor):
        static.append(i + 1)  # weight is 0
    static_argnames = []
    for k, v in kwargs.items():
      if not isinstance(v, torch.Tensor):
        static_argnames.append(k)

    def f(weights, *args, **kwargs):
      weights, args, kwargs = torchax.tensor.wrap((weights, args, kwargs))
      with torchax.functions.XLAFunctionMode(), torchax.tensor.XLADispatchMode():
        res = torch.func.functional_call(self._model, weights, args, kwargs)
        if isinstance(res, tuple) and len(res) == 1:
          res = res[0]
      return torchax.tensor.unwrap(res)

    fjit = jax.jit(f, static_argnames=tuple(static_argnames))
    return torch_view(fjit)

  def forward(self, *args, **kwargs):
    (args, kwargs) = pytree.tree_map(self._maybe_move_tensor, (args, kwargs))
    if self._func_jitted_torch is None:
      self._func_jitted_torch = self._make_jitted(args, kwargs)
    return self._func_jitted_torch(self._weights, *args, **kwargs)

  def __call__(self, *args, **kwargs):
    return self.forward(*args, **kwargs)

  def __getattr__(self, key):
    return getattr(self._model, key)


def compile_pipe(pipe):
  pipe.text_encoder = CompiledModule(pipe.text_encoder)
  pipe.text_encoder_2 = CompiledModule(pipe.text_encoder_2)
  pipe.unet = CompiledModule(pipe.unet)
  pipe.vae = CompiledModule(pipe.vae)


def main():
  pipe = DiffusionPipeline.from_pretrained(
    # "stabilityai/stable-diffusion-xl-base-0.9",
    "stabilityai/stable-diffusion-xl-base-1.0",
    use_safetensors=True,
  )
  compile_pipe(pipe)

  global_bs = 10
  inference_steps = 20
  resol = 1024
  prompts = ["a photo of an astronaut riding a horse on mars"] * global_bs
  print(
    f"global batch size {global_bs}",
    f"inference steps {inference_steps}",
    f"Image resolution {resol}",
    flush=True,
  )

  iters = 5
  for i in range(iters):
    prompt = prompts
    # print('per device prompts len',len(prompt))
    # prompt = prompts[rank]
    start = time()
    image = pipe(
      prompt, num_inference_steps=inference_steps, height=resol, width=resol
    ).images[0]
    print(f"Step {i} inference time {time() - start} sec", flush=True)


if __name__ == "__main__":
  main()