fix: dtype alignment for attention mask and add bf16 support

.pre-commit-config.yaml

@@ -21,6 +21,7 @@ repos:
     hooks:
       - id: pretty-format-toml
         args: [--autofix]
+        exclude: ^uv\.lock$
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.13.1

cccv/arch/sr/dat_arch.py

@@ -425,14 +425,14 @@ def forward(self, qkv, H, W, mask=None):
                 self.H_sp * self.W_sp, self.H_sp * self.W_sp, -1
             )
             relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
-            attn = attn + relative_position_bias.unsqueeze(0)
+            attn = attn + relative_position_bias.unsqueeze(0).to(attn.dtype)
 
         N = attn.shape[3]
 
         # use mask for shift window
         if mask is not None:
             nW = mask.shape[0]
-            attn = attn.view(B, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(B, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0).to(attn.dtype)
             attn = attn.view(-1, self.num_heads, N, N)
 
         attn = nn.functional.softmax(attn, dim=-1, dtype=attn.dtype)

cccv/arch/sr/hat_arch.py

@@ -477,11 +477,11 @@ def forward(self, x, rpi, mask=None):
             -1,
         )  # Wh*Ww,Wh*Ww,nH
         relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
-        attn = attn + relative_position_bias.unsqueeze(0)
+        attn = attn + relative_position_bias.unsqueeze(0).to(attn.dtype)
 
         if mask is not None:
             nw = mask.shape[0]
-            attn = attn.view(b_ // nw, nw, self.num_heads, n, n) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(b_ // nw, nw, self.num_heads, n, n) + mask.unsqueeze(1).unsqueeze(0).to(attn.dtype)
             attn = attn.view(-1, self.num_heads, n, n)
             attn = self.softmax(attn)
         else:
@@ -750,7 +750,7 @@ def forward(self, x, x_size, rpi):
             -1,
         )  # ws*ws, wse*wse, nH
         relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, ws*ws, wse*wse
-        attn = attn + relative_position_bias.unsqueeze(0)
+        attn = attn + relative_position_bias.unsqueeze(0).to(attn.dtype)
 
         attn = self.softmax(attn)
         attn_windows = (attn @ v).transpose(1, 2).reshape(b_, nq, self.dim)

cccv/arch/sr/swinir_arch.py

@@ -402,11 +402,11 @@ def forward(self, x, mask=None):
             self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1
         )  # Wh*Ww,Wh*Ww,nH
         relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()  # nH, Wh*Ww, Wh*Ww
-        attn = attn + relative_position_bias.unsqueeze(0)
+        attn = attn + relative_position_bias.unsqueeze(0).to(attn.dtype)
 
         if mask is not None:
             nW = mask.shape[0]
-            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0)
+            attn = attn.view(B_ // nW, nW, self.num_heads, N, N) + mask.unsqueeze(1).unsqueeze(0).to(attn.dtype)
             attn = attn.view(-1, self.num_heads, N, N)
             attn = self.softmax(attn)
         else:

cccv/auto/model.py

@@ -16,11 +16,14 @@ def from_pretrained(
         *,
         device: Optional[torch.device] = None,
         fp16: bool = True,
+        bf16: bool = False,
         compile: bool = False,
         compile_backend: Optional[str] = None,
         tile: Optional[Tuple[int, int]] = (128, 128),
         tile_pad: int = 8,
         pad_img: Optional[Tuple[int, int]] = None,
+        bf16_preflight: bool = True,
+        bf16_preflight_size: Tuple[int, int] = (64, 64),
         model_dir: Optional[Union[Path, str]] = None,
         gh_proxy: Optional[str] = None,
         **kwargs: Any,
@@ -30,12 +33,15 @@ def from_pretrained(
 
         :param pretrained_model_name_or_path:
         :param device: inference device
-        :param fp16: use fp16 precision or not
+        :param fp16: use fp16 (half) precision or not
+        :param bf16: use bf16 (bfloat16) precision or not, takes precedence over fp16
         :param compile: use torch.compile or not
         :param compile_backend: backend of torch.compile
         :param tile: tile size for tile inference, tile[0] is width, tile[1] is height, None for disable
         :param tile_pad: The padding size for each tile
         :param pad_img: The size for the padded image, pad[0] is width, pad[1] is height, None for auto calculate
+        :param bf16_preflight: run a small bf16 inference before actual user inference, fallback if it fails
+        :param bf16_preflight_size: The bf16 preflight input size as (height, width), aligned per model if needed
         :param model_dir: The path to cache the downloaded model. Should be a full path. If None, use default cache path.
         :param gh_proxy: The proxy for downloading from github release. Example: https://github.abskoop.workers.dev/
         :return:
@@ -46,11 +52,14 @@ def from_pretrained(
             config=config,
             device=device,
             fp16=fp16,
+            bf16=bf16,
             compile=compile,
             compile_backend=compile_backend,
             tile=tile,
             tile_pad=tile_pad,
             pad_img=pad_img,
+            bf16_preflight=bf16_preflight,
+            bf16_preflight_size=bf16_preflight_size,
             model_dir=model_dir,
             gh_proxy=gh_proxy,
             **kwargs,
@@ -62,11 +71,14 @@ def from_config(
         *,
         device: Optional[torch.device] = None,
         fp16: bool = True,
+        bf16: bool = False,
         compile: bool = False,
         compile_backend: Optional[str] = None,
         tile: Optional[Tuple[int, int]] = (128, 128),
         tile_pad: int = 8,
         pad_img: Optional[Tuple[int, int]] = None,
+        bf16_preflight: bool = True,
+        bf16_preflight_size: Tuple[int, int] = (64, 64),
         model_dir: Optional[Union[Path, str]] = None,
         gh_proxy: Optional[str] = None,
         **kwargs: Any,
@@ -76,12 +88,15 @@ def from_config(
 
         :param config: The config object. We suggest use cccv.BaseConfig or its subclass.
         :param device: inference device
-        :param fp16: use fp16 precision or not
+        :param fp16: use fp16 (half) precision or not
+        :param bf16: use bf16 (bfloat16) precision or not, takes precedence over fp16
         :param compile: use torch.compile or not
         :param compile_backend: backend of torch.compile
         :param tile: tile size for tile inference, tile[0] is width, tile[1] is height, None for disable
         :param tile_pad: The padding size for each tile
         :param pad_img: The size for the padded image, pad[0] is width, pad[1] is height, None for auto calculate
+        :param bf16_preflight: run a small bf16 inference before actual user inference, fallback if it fails
+        :param bf16_preflight_size: The bf16 preflight input size as (height, width), aligned per model if needed
         :param model_dir: The path to cache the downloaded model. Should be a full path. If None, use default cache path.
         :param gh_proxy: The proxy for downloading from github release. Example: https://github.abskoop.workers.dev/
         :return:
@@ -91,11 +106,14 @@ def from_config(
             config=config,
             device=device,
             fp16=fp16,
+            bf16=bf16,
             compile=compile,
             compile_backend=compile_backend,
             tile=tile,
             tile_pad=tile_pad,
             pad_img=pad_img,
+            bf16_preflight=bf16_preflight,
+            bf16_preflight_size=bf16_preflight_size,
             model_dir=model_dir,
             gh_proxy=gh_proxy,
             **kwargs,

cccv/model/base_model.py

@@ -2,7 +2,7 @@
 import warnings
 from inspect import signature
 from pathlib import Path
-from typing import Any, Optional, Tuple, Union
+from typing import Any, Dict, Optional, Tuple, Union
 
 import torch
 
@@ -19,12 +19,15 @@ class CCBaseModel(BaseModelInterface):
 
     :param config: config of the model
     :param device: inference device
-    :param fp16: use fp16 precision or not
+    :param fp16: use fp16 (half) precision or not
+    :param bf16: use bf16 (bfloat16) precision or not, takes precedence over fp16; wider dynamic range than fp16 to avoid overflow/NaN on some models (e.g. transformer-based SR)
     :param compile: use torch.compile or not
     :param compile_backend: backend of torch.compile
     :param tile: tile size for tile inference, tile[0] is width, tile[1] is height, None for disable
     :param tile_pad: The padding size for each tile
     :param pad_img: The size for the padded image, pad[0] is width, pad[1] is height, None for auto calculate
+    :param bf16_preflight: run a small bf16 inference before actual user inference, fallback if it fails
+    :param bf16_preflight_size: The bf16 preflight input size as (height, width), aligned per model if needed
     :param model_dir: The path to cache the downloaded model. Should be a full path. If None, use default cache path.
     :param gh_proxy: The proxy for downloading from github release. Example: https://github.abskoop.workers.dev/
     """
@@ -34,11 +37,14 @@ def __init__(
         config: Any,
         device: Optional[torch.device] = None,
         fp16: bool = True,
+        bf16: bool = False,
         compile: bool = False,
         compile_backend: Optional[str] = None,
         tile: Optional[Tuple[int, int]] = (128, 128),
         tile_pad: int = 8,
         pad_img: Optional[Tuple[int, int]] = None,
+        bf16_preflight: bool = True,
+        bf16_preflight_size: Tuple[int, int] = (64, 64),
         model_dir: Optional[Union[Path, str]] = None,
         gh_proxy: Optional[str] = None,
         **kwargs: Any,
@@ -53,11 +59,16 @@ def __init__(
         self.config = config
         self.device: Optional[torch.device] = device
         self.fp16: bool = fp16
+        self.bf16: bool = bf16
+        # half precision dtype, bf16 takes precedence over fp16
+        self.half_dtype: torch.dtype = torch.bfloat16 if bf16 else torch.float16
         self.compile: bool = compile
         self.compile_backend: Optional[str] = compile_backend
         self.tile: Optional[Tuple[int, int]] = tile
         self.tile_pad: int = tile_pad
         self.pad_img: Optional[Tuple[int, int]] = pad_img
+        self.bf16_preflight: bool = bf16_preflight
+        self.bf16_preflight_size: Tuple[int, int] = bf16_preflight_size
         self.model_dir: Optional[Union[Path, str]] = model_dir
         self.gh_proxy: Optional[str] = gh_proxy
 
@@ -69,26 +80,140 @@ def __init__(
 
         self.model: torch.nn.Module = self.load_model()
 
-        # fp16
-        if self.fp16:
+        # half precision (fp16 or bf16, bf16 takes precedence)
+        if self.fp16 or self.bf16:
+            self._try_enable_half_precision()
+
+        # compile
+        if self.compile:
+            self._try_compile_model()
+
+        if self.bf16 and self.bf16_preflight:
+            self._run_bf16_preflight_or_fallback()
+
+    def _try_enable_half_precision(self) -> None:
+        self.half_dtype = torch.bfloat16 if self.bf16 else torch.float16
+
+        try:
+            self.model = self.model.to(self.half_dtype)
+        except Exception as e:
+            if self.bf16 and self.fp16:
+                warnings.warn(f"[CCCV] {e}. bf16 is not supported on this model, fallback to fp16.", stacklevel=2)
+                self.bf16 = False
+                self.half_dtype = torch.float16
+                self.model = self.load_model()
+                try:
+                    self.model = self.model.to(self.half_dtype)
+                    return
+                except Exception as fp16_error:
+                    warnings.warn(f"[CCCV] {fp16_error}. fp16 fallback failed, fallback to fp32.", stacklevel=2)
+                    self.fp16 = False
+                    self.half_dtype = torch.float32
+                    self.model = self.load_model()
+                    return
+
+            warnings.warn(f"[CCCV] {e}. half precision is not supported on this model, fallback to fp32.", stacklevel=2)
+            self.bf16 = False
+            self.fp16 = False
+            self.half_dtype = torch.float32
+            self.model = self.load_model()
+
+    def _try_compile_model(self) -> None:
+        try:
+            if self.compile_backend is None:
+                if sys.platform == "darwin":
+                    self.compile_backend = "aot_eager"
+                else:
+                    self.compile_backend = "inductor"
+            self.model = torch.compile(self.model, backend=self.compile_backend)
+        except Exception as e:
+            warnings.warn(f"[CCCV] {e}, compile is not supported on this model.", stacklevel=2)
+
+    def _run_bf16_preflight_or_fallback(self) -> None:
+        preflight = None
+        try:
+            preflight = self.get_bf16_preflight_inputs()
+            if preflight is None:
+                return
+
+            preflight_args, preflight_kwargs = preflight
+            with torch.inference_mode():
+                out = self.inference(*preflight_args, **preflight_kwargs)
+            self._synchronize_device()
+            del out
+        except Exception as e:
+            warnings.warn(
+                f"[CCCV] {e}. bf16 inference is not supported on this device, disabling bf16 as fallback.",
+                stacklevel=2,
+            )
+            self._fallback_from_bf16()
+        finally:
+            preflight = None
+            self._empty_device_cache()
+
+    def _fallback_from_bf16(self) -> None:
+        self.bf16 = False
+        self.half_dtype = torch.float16 if self.fp16 else torch.float32
+        self.model = self.load_model()
+
+        if self.fp16 or self.bf16:
             try:
-                self.model = self.model.half()
+                self.model = self.model.to(self.half_dtype)
             except Exception as e:
-                warnings.warn(f"[CCCV] {e}. fp16 is not supported on this model, fallback to fp32.", stacklevel=2)
+                warnings.warn(f"[CCCV] {e}. fp16 fallback failed, fallback to fp32.", stacklevel=2)
                 self.fp16 = False
+                self.half_dtype = torch.float32
                 self.model = self.load_model()
 
-        # compile
         if self.compile:
-            try:
-                if self.compile_backend is None:
-                    if sys.platform == "darwin":
-                        self.compile_backend = "aot_eager"
-                    else:
-                        self.compile_backend = "inductor"
-                self.model = torch.compile(self.model, backend=self.compile_backend)
-            except Exception as e:
-                warnings.warn(f"[CCCV] {e}, compile is not supported on this model.", stacklevel=2)
+            self._try_compile_model()
+
+    def _synchronize_device(self) -> None:
+        if self.device is not None and torch.device(self.device).type == "cuda" and torch.cuda.is_available():
+            torch.cuda.synchronize(torch.device(self.device))
+
+    def _empty_device_cache(self) -> None:
+        if self.device is not None and torch.device(self.device).type == "cuda" and torch.cuda.is_available():
+            torch.cuda.empty_cache()
+
+    @staticmethod
+    def _tensor_to_numpy(tensor: torch.Tensor) -> Any:
+        tensor = tensor.cpu()
+        if tensor.dtype == torch.bfloat16:
+            tensor = tensor.float()
+        return tensor.numpy()
+
+    def get_bf16_preflight_inputs(self) -> Optional[Tuple[Tuple[Any, ...], Dict[str, Any]]]:
+        """
+        Hook: Subclasses can return representative inference args/kwargs for bf16 device validation.
+
+        By default, no preflight runs because the base model does not know the input contract.
+        """
+        return None
+
+    def _get_bf16_preflight_image_size(self, multiple: int = 1) -> Tuple[int, int]:
+        height, width = self.bf16_preflight_size
+        input_multiple = max(1, multiple, self._infer_bf16_preflight_multiple())
+        return self._ceil_to_multiple(height, input_multiple), self._ceil_to_multiple(width, input_multiple)
+
+    def _infer_bf16_preflight_multiple(self) -> int:
+        multiples = []
+
+        window_size = getattr(self.config, "window_size", None)
+        if isinstance(window_size, int):
+            multiples.append(window_size)
+
+        split_size = getattr(self.config, "split_size", None)
+        if isinstance(split_size, int):
+            multiples.append(split_size)
+        elif isinstance(split_size, (list, tuple)):
+            multiples.extend(value for value in split_size if isinstance(value, int))
+
+        return max(multiples) if multiples else 1
+
+    @staticmethod
+    def _ceil_to_multiple(value: int, multiple: int) -> int:
+        return ((value + multiple - 1) // multiple) * multiple
 
     def post_init_hook(self) -> None:
         """