[TPU] Initial TorchTPU backend integration (eager + torch.compile)

docs/source/en/optimization/tpu.md

@@ -0,0 +1,123 @@
+<!--Copyright 2026 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
+
+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.
+-->
+
+# TorchTPU
+
+[TorchTPU](https://github.com/google-pytorch/torch_tpu/) provides a PyTorch backend for Google's Tensor Processing Units (TPUs), enabling you to run diffusers pipelines on Google Cloud TPUs (v6e, v5p, …) with minimal code changes.
+
+Four execution modes are available:
+
+| Mode | Constant | How to activate | Notes |
+|---|---|---|---|
+| **Strict Eager** (default) | `EagerMode.DEFER_NEVER` | just `import torch_tpu` | Operations dispatched one at a time, asynchronous |
+| **Compile** | — | `pipe.enable_tpu_compile()` | AOT compilation with `TpuBackend` |
+
+## Installation
+
+Follow the [TorchTPU installation guide](https://github.com/google-pytorch/torch_tpu/). After installation,
+`import torch_tpu` registers the `"tpu"` device automatically.
+
+## Basic usage (strict eager mode)
+
+```python
+import torch
+import torch_tpu  # noqa: F401 — registers torch.tpu
+
+from diffusers import FluxPipeline
+
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell",
+    torch_dtype=torch.bfloat16,
+)
+
+# Move only the denoising components to TPU; text encoders stay on CPU.
+pipe.transformer.to("tpu")
+pipe.vae.to("tpu")
+
+# _execution_device is now "tpu" automatically.
+image = pipe(
+    prompt="a golden retriever surfing a wave, photorealistic",
+    height=1024,
+    width=1024,
+    num_inference_steps=4,
+    guidance_scale=0.0,
+).images[0]
+
+image.save("output.png")
+```
+
+## Compiled mode (recommended for production)
+
+`torch.compile` with `TpuBackend` traces the transformer statically. The first call (warmup)
+is slow because it triggers compilation; subsequent calls reuse the compiled graph.
+
+> [!IMPORTANT]
+> TorchTPU requires **static shapes** — `torch.compile` is called with `dynamic=False`
+> internally. Every time `height`, `width`, or `num_inference_steps` changes, the graph is
+> recompiled from scratch. Keep these values constant across all calls after warmup, or call
+> `tpu_warmup` again before changing them.
+
+```python
+import torch
+import torch_tpu  # noqa: F401
+
+from diffusers import FluxPipeline
+
+pipe = FluxPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-schnell",
+    torch_dtype=torch.bfloat16,
+)
+pipe.transformer.to("tpu")
+pipe.vae.to("tpu")
+
+# Compile TPU components with TpuBackend.
+# Also applies AttnProcessor to replace SDP-based attention (required for XLA).
+pipe.enable_tpu_compile()
+
+# Warmup — triggers static graph compilation.
+pipe.tpu_warmup(
+    prompt="warmup",
+    height=1024,
+    width=1024,
+    num_inference_steps=4,
+    guidance_scale=0.0,
+)
+
+# Timed inference reuses the compiled graph.
+image = pipe(
+    prompt="a golden retriever surfing a wave, photorealistic",
+    height=1024,
+    width=1024,
+    num_inference_steps=4,
+    guidance_scale=0.0,
+).images[0]
+
+image.save("output.png")
+```
+
+## Eager mode
+
+TorchTPU defaults to **Strict Eager** (`EagerMode.DEFER_NEVER`): operations are dispatched one
+at a time asynchronously, matching standard PyTorch GPU behaviour.
+
+> [!TIP]
+> For the best production throughput, prefer `torch.compile` via `pipe.enable_tpu_compile()`.
+
+## API reference
+
+### `enable_tpu_compile`
+
+[[autodoc]] diffusers.DiffusionPipeline.enable_tpu_compile
+
+### `tpu_warmup`
+
+[[autodoc]] diffusers.DiffusionPipeline.tpu_warmup

src/diffusers/models/unets/unet_2d_condition.py

@@ -864,7 +864,14 @@ def get_time_embed(self, sample: torch.Tensor, timestep: torch.Tensor | float |
         # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
         timesteps = timesteps.expand(sample.shape[0])
 
-        t_emb = self.time_proj(timesteps)
+        # On TPU in eager/lazy mode, torch.cat([sin, cos], dim=-1) inside time_proj
+        # lands at an unaligned offset in the XLA DUS fusion emitter → crash.
+        # torch.compile with TpuBackend handles this internally, so skip the CPU
+        # workaround when we're inside a compiled graph.
+        if sample.device.type == "tpu" and not torch.compiler.is_compiling():
+            t_emb = self.time_proj(timesteps.cpu()).to(sample.device)
+        else:
+            t_emb = self.time_proj(timesteps)
         # `Timesteps` does not contain any weights and will always return f32 tensors
         # but time_embedding might actually be running in fp16. so we need to cast here.
         # there might be better ways to encapsulate this.

src/diffusers/pipelines/ernie_image/pipeline_ernie_image.py

@@ -114,7 +114,7 @@ def _enhance_prompt_with_pe(
             tokenize=False,
             add_generation_prompt=False,  # "Output:" is already in the user block
         )
-        inputs = self.pe_tokenizer(input_text, return_tensors="pt").to(device)
+        inputs = self.pe_tokenizer(input_text, return_tensors="pt").to(self.pe.device)
         output_ids = self.pe.generate(
             **inputs,
             max_new_tokens=self.pe_tokenizer.model_max_length,
@@ -155,7 +155,7 @@ def encode_prompt(
                 else:
                     ids = [0]
 
-            input_ids = torch.tensor([ids], device=device)
+            input_ids = torch.tensor([ids], device=self.text_encoder.device)
             with torch.no_grad():
                 outputs = self.text_encoder(
                     input_ids=input_ids,

src/diffusers/pipelines/flux/pipeline_flux.py

@@ -250,7 +250,8 @@ def _get_t5_prompt_embeds(
                 f" {max_sequence_length} tokens: {removed_text}"
             )
 
-        prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0]
+        model_device = self.text_encoder_2.device
+        prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0]
 
         dtype = self.text_encoder_2.dtype
         prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
@@ -295,7 +296,8 @@ def _get_clip_prompt_embeds(
                 "The following part of your input was truncated because CLIP can only handle sequences up to"
                 f" {self.tokenizer_max_length} tokens: {removed_text}"
             )
-        prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False)
+        model_device = self.text_encoder.device
+        prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False)
 
         # Use pooled output of CLIPTextModel
         prompt_embeds = prompt_embeds.pooler_output

src/diffusers/pipelines/flux/pipeline_flux_controlnet.py

@@ -282,7 +282,8 @@ def _get_t5_prompt_embeds(
                 f" {max_sequence_length} tokens: {removed_text}"
             )
 
-        prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0]
+        model_device = self.text_encoder_2.device
+        prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0]
 
         dtype = self.text_encoder_2.dtype
         prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
@@ -327,7 +328,8 @@ def _get_clip_prompt_embeds(
                 "The following part of your input was truncated because CLIP can only handle sequences up to"
                 f" {self.tokenizer_max_length} tokens: {removed_text}"
             )
-        prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False)
+        model_device = self.text_encoder.device
+        prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False)
 
         # Use pooled output of CLIPTextModel
         prompt_embeds = prompt_embeds.pooler_output

src/diffusers/pipelines/flux/pipeline_flux_inpaint.py

@@ -275,7 +275,8 @@ def _get_t5_prompt_embeds(
                 f" {max_sequence_length} tokens: {removed_text}"
             )
 
-        prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0]
+        model_device = self.text_encoder_2.device
+        prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0]
 
         dtype = self.text_encoder_2.dtype
         prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
@@ -321,7 +322,8 @@ def _get_clip_prompt_embeds(
                 "The following part of your input was truncated because CLIP can only handle sequences up to"
                 f" {self.tokenizer_max_length} tokens: {removed_text}"
             )
-        prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False)
+        model_device = self.text_encoder.device
+        prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False)
 
         # Use pooled output of CLIPTextModel
         prompt_embeds = prompt_embeds.pooler_output

src/diffusers/pipelines/flux/pipeline_flux_kontext.py

@@ -296,7 +296,8 @@ def _get_t5_prompt_embeds(
                 f" {max_sequence_length} tokens: {removed_text}"
             )
 
-        prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0]
+        model_device = self.text_encoder_2.device
+        prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0]
 
         dtype = self.text_encoder_2.dtype
         prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
@@ -342,7 +343,8 @@ def _get_clip_prompt_embeds(
                 "The following part of your input was truncated because CLIP can only handle sequences up to"
                 f" {self.tokenizer_max_length} tokens: {removed_text}"
             )
-        prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False)
+        model_device = self.text_encoder.device
+        prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False)
 
         # Use pooled output of CLIPTextModel
         prompt_embeds = prompt_embeds.pooler_output

src/diffusers/pipelines/flux/pipeline_flux_kontext_inpaint.py

@@ -329,7 +329,8 @@ def _get_t5_prompt_embeds(
                 f" {max_sequence_length} tokens: {removed_text}"
             )
 
-        prompt_embeds = self.text_encoder_2(text_input_ids.to(device), output_hidden_states=False)[0]
+        model_device = self.text_encoder_2.device
+        prompt_embeds = self.text_encoder_2(text_input_ids.to(model_device), output_hidden_states=False)[0]
 
         dtype = self.text_encoder_2.dtype
         prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
@@ -375,7 +376,8 @@ def _get_clip_prompt_embeds(
                 "The following part of your input was truncated because CLIP can only handle sequences up to"
                 f" {self.tokenizer_max_length} tokens: {removed_text}"
             )
-        prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False)
+        model_device = self.text_encoder.device
+        prompt_embeds = self.text_encoder(text_input_ids.to(model_device), output_hidden_states=False)
 
         # Use pooled output of CLIPTextModel
         prompt_embeds = prompt_embeds.pooler_output

src/diffusers/pipelines/flux2/pipeline_flux2_klein.py

@@ -241,8 +241,9 @@ def _get_qwen3_prompt_embeds(
             all_input_ids.append(inputs["input_ids"])
             all_attention_masks.append(inputs["attention_mask"])
 
-        input_ids = torch.cat(all_input_ids, dim=0).to(device)
-        attention_mask = torch.cat(all_attention_masks, dim=0).to(device)
+        model_device = text_encoder.device
+        input_ids = torch.cat(all_input_ids, dim=0).to(model_device)
+        attention_mask = torch.cat(all_attention_masks, dim=0).to(model_device)
 
         # Forward pass through the model
         output = text_encoder(

src/diffusers/pipelines/pipeline_utils.py

@@ -22,7 +22,7 @@
 import types
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Union, get_args, get_origin, get_type_hints
+from typing import Any, Callable, Dict, List, Optional, Union, get_args, get_origin, get_type_hints
 
 import httpx
 import numpy as np
@@ -73,6 +73,7 @@
     is_transformers_version,
     logging,
     numpy_to_pil,
+    requires_backends,
 )
 from ..utils.distributed_utils import is_torch_dist_rank_zero
 from ..utils.hub_utils import (
@@ -1147,6 +1148,15 @@ def _execution_device(self):
             except ValueError:
                 pass
 
+        # When text encoders are offloaded to CPU while the denoising backbone
+        # (unet, transformer, vae) runs on an accelerator, self.device returns CPU
+        # (first component). Prefer any non-CPU, non-meta component device so that
+        # scheduler and latent tensors land on the accelerator. This covers TPU,
+        # NPU (npu), Intel GPU (xpu), Habana (hpu), and any other backend.
+        for name, model in self.components.items():
+            if isinstance(model, torch.nn.Module) and model.device.type not in ("cpu", "meta"):
+                return model.device
+
         for name, model in self.components.items():
             if not isinstance(model, torch.nn.Module) or name in self._exclude_from_cpu_offload:
                 continue
@@ -2339,3 +2349,89 @@ def unfuse_qkv_projections(self, unet: bool = True, vae: bool = True):
             else:
                 self.vae.unfuse_qkv_projections()
                 self.fusing_vae = False
+
+    def enable_tpu_compile(
+        self,
+        model_names: Optional[List[str]] = None,
+        **compile_kwargs,
+    ) -> None:
+        """Compile pipeline components that are on TPU using ``torch.compile`` with the ``TpuBackend``.
+
+        Before compiling, each component that exposes ``set_attn_processor`` has ``AttnProcessor``
+        applied. This replaces ``AttnProcessor2_0`` (SDP-based) which triggers XLA fusion-emitter
+        crashes in eager/lazy mode. ``TpuBackend`` handles the resulting ``torch.cat`` layout
+        internally during static tracing, so no additional wrapper is needed at compile time.
+
+        Args:
+            model_names (`list[str]`, *optional*):
+                Names of pipeline components to compile. Defaults to all ``torch.nn.Module``
+                components currently resident on a TPU device.
+            **compile_kwargs:
+                Extra keyword arguments forwarded to ``torch.compile``. ``backend`` defaults to
+                ``TpuBackend()`` and ``dynamic`` defaults to ``False`` (required for static tracing).
+
+        Example:
+        ```python
+        import torch
+        import torch_tpu  # noqa: F401
+
+        pipe.transformer.to("tpu")
+        pipe.vae.to("tpu")
+        pipe.enable_tpu_compile()
+        ```
+        """
+        requires_backends(self, "torch_tpu")
+        from torch_tpu._internal.compile import TpuBackend
+
+        from ..models.attention_processor import AttnProcessor
+
+        if model_names is None:
+            model_names = [
+                name
+                for name, comp in self.components.items()
+                if isinstance(comp, torch.nn.Module) and comp.device.type == "tpu"
+            ]
+
+        for name in model_names:
+            component = getattr(self, name, None)
+            if not isinstance(component, torch.nn.Module):
+                logger.warning(f"`enable_tpu_compile`: component '{name}' is not a nn.Module, skipping.")
+                continue
+            if is_compiled_module(component):
+                logger.warning(f"`enable_tpu_compile`: component '{name}' is already compiled, skipping.")
+                continue
+            if hasattr(component, "set_attn_processor"):
+                component.set_attn_processor(AttnProcessor())
+            compile_kwargs.setdefault("backend", TpuBackend())
+            compile_kwargs.setdefault("dynamic", False)
+            logger.info(f"Compiling '{name}' with TpuBackend.")
+            setattr(self, name, torch.compile(component, **compile_kwargs))
+
+    def tpu_warmup(self, *args, **kwargs) -> None:
+        """Run a single forward pass to trigger XLA / ``TpuBackend`` compilation.
+
+        Call this after ``enable_tpu_compile`` and before timed inference. The warmup
+        pass compiles the static computation graphs; subsequent calls reuse the compiled
+        graphs and run at full speed.
+
+        Args:
+            *args: Positional arguments forwarded to the pipeline ``__call__``.
+            **kwargs: Keyword arguments forwarded to the pipeline ``__call__``.
+
+        Example:
+        ```python
+        pipe.tpu_warmup(
+            prompt="warmup",
+            height=1024,
+            width=1024,
+            num_inference_steps=4,
+            guidance_scale=0.0,
+        )
+        ```
+        """
+        logger.info("Running TPU warmup pass to trigger XLA compilation...")
+        with torch.no_grad():
+            self(*args, **kwargs)
+        if hasattr(torch, "tpu") and hasattr(torch.tpu, "synchronize"):
+            torch.tpu.synchronize()
+        logger.info("TPU warmup complete.")