diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d529b19..be445f9 100644 --- a/.pre-commit-config.yaml +++ b/.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 diff --git a/cccv/arch/sr/dat_arch.py b/cccv/arch/sr/dat_arch.py index 5eec59f..7e9e7cd 100644 --- a/cccv/arch/sr/dat_arch.py +++ b/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) diff --git a/cccv/arch/sr/hat_arch.py b/cccv/arch/sr/hat_arch.py index b90ce7b..6c78714 100644 --- a/cccv/arch/sr/hat_arch.py +++ b/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) diff --git a/cccv/arch/sr/swinir_arch.py b/cccv/arch/sr/swinir_arch.py index e3f7115..47612b6 100644 --- a/cccv/arch/sr/swinir_arch.py +++ b/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: diff --git a/cccv/auto/model.py b/cccv/auto/model.py index 472d0da..f20a894 100644 --- a/cccv/auto/model.py +++ b/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, diff --git a/cccv/model/base_model.py b/cccv/model/base_model.py index 9c879eb..40f45a8 100644 --- a/cccv/model/base_model.py +++ b/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: """ diff --git a/cccv/model/sr_base_model.py b/cccv/model/sr_base_model.py index 7bb9242..b1d1a11 100644 --- a/cccv/model/sr_base_model.py +++ b/cccv/model/sr_base_model.py @@ -1,4 +1,4 @@ -from typing import Any +from typing import Any, Dict, Optional, Tuple import cv2 import numpy as np @@ -14,6 +14,11 @@ @MODEL_REGISTRY.register(name=ModelType.SRBaseModel) class SRBaseModel(CCBaseModel): + def get_bf16_preflight_inputs(self) -> Optional[Tuple[Tuple[Any, ...], Dict[str, Any]]]: + height, width = self._get_bf16_preflight_image_size() + img = torch.zeros((1, 3, height, width), device=self.device, dtype=self.half_dtype) + return (img,), {} + @torch.inference_mode() # type: ignore def inference(self, img: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: cfg: SRBaseConfig = self.config @@ -42,11 +47,11 @@ def inference_image(self, img: np.ndarray, *args: Any, **kwargs: Any) -> np.ndar img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) img = transforms.ToTensor()(img).unsqueeze(0).to(self.device) - if self.fp16: - img = img.half() + if self.fp16 or self.bf16: + img = img.to(self.half_dtype) img = self.inference(img) - img = img.squeeze(0).permute(1, 2, 0).cpu().numpy() + img = self._tensor_to_numpy(img.squeeze(0).permute(1, 2, 0)) img = (img * 255).clip(0, 255).astype("uint8") img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) diff --git a/cccv/model/vfi/drba_model.py b/cccv/model/vfi/drba_model.py index 681776a..f13407d 100644 --- a/cccv/model/vfi/drba_model.py +++ b/cccv/model/vfi/drba_model.py @@ -1,4 +1,4 @@ -from typing import Any, List +from typing import Any, Dict, List, Optional, Tuple import cv2 import numpy as np @@ -17,6 +17,11 @@ class DRBAModel(VFIBaseModel): def post_init_hook(self) -> None: self.load_state_dict_strict = False + def get_bf16_preflight_inputs(self) -> Optional[Tuple[Tuple[Any, ...], Dict[str, Any]]]: + height, width = self._get_bf16_preflight_image_size(multiple=32) + imgs = torch.zeros((1, 3, 3, height, width), device=self.device, dtype=self.half_dtype) + return (imgs, [-1, -0.5], [0], [0.5, 1], False, False, 1.0, None), {} + def transform_state_dict(self, state_dict: Any) -> Any: def _convert(param: Any) -> Any: return {k.replace("module.", ""): v for k, v in param.items() if "module." in k} @@ -86,15 +91,15 @@ def inference_image_list(self, img_list: List[np.ndarray], *args: Any, **kwargs: # b, n, c, h, w img_tensor_stack = torch.stack(new_img_list, dim=1) - if self.fp16: - img_tensor_stack = img_tensor_stack.half() + if self.fp16 or self.bf16: + img_tensor_stack = img_tensor_stack.to(self.half_dtype) results, _ = self.inference(img_tensor_stack, [-1, -0.5], [0], [0.5, 1], False, False, 1.0, None) results_list = [] for i in range(results.shape[1]): img = results[0, i, :, :, :] - img = img.permute(1, 2, 0).cpu().numpy() + img = self._tensor_to_numpy(img.permute(1, 2, 0)) img = (img * 255).clip(0, 255).astype("uint8") img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) results_list.append(img) diff --git a/cccv/model/vfi/rife_model.py b/cccv/model/vfi/rife_model.py index 0e8605e..9bafa46 100644 --- a/cccv/model/vfi/rife_model.py +++ b/cccv/model/vfi/rife_model.py @@ -1,4 +1,4 @@ -from typing import Any, List +from typing import Any, Dict, List, Optional, Tuple import cv2 import numpy as np @@ -16,6 +16,11 @@ class RIFEModel(VFIBaseModel): def post_init_hook(self) -> None: self.load_state_dict_strict = False + def get_bf16_preflight_inputs(self) -> Optional[Tuple[Tuple[Any, ...], Dict[str, Any]]]: + height, width = self._get_bf16_preflight_image_size(multiple=32) + imgs = torch.zeros((1, 2, 3, height, width), device=self.device, dtype=self.half_dtype) + return (imgs, 0.5, 1.0), {} + def transform_state_dict(self, state_dict: Any) -> Any: def _convert(param: Any) -> Any: return {k.replace("module.", ""): v for k, v in param.items() if "module." in k} @@ -68,15 +73,15 @@ def inference_image_list(self, img_list: List[np.ndarray], *args: Any, **kwargs: # b, n, c, h, w img_tensor_stack = torch.stack(new_img_list, dim=1) - if self.fp16: - img_tensor_stack = img_tensor_stack.half() + if self.fp16 or self.bf16: + img_tensor_stack = img_tensor_stack.to(self.half_dtype) out = self.inference(img_tensor_stack, timestep=0.5, scale=1.0) # Convert to numpy image list results_list = [] - img = out.squeeze(0).permute(1, 2, 0).cpu().numpy() + img = self._tensor_to_numpy(out.squeeze(0).permute(1, 2, 0)) img = (img * 255).clip(0, 255).astype("uint8") img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) diff --git a/cccv/model/vsr_base_model.py b/cccv/model/vsr_base_model.py index 487e5a0..e53e5b0 100644 --- a/cccv/model/vsr_base_model.py +++ b/cccv/model/vsr_base_model.py @@ -1,4 +1,4 @@ -from typing import Any, List +from typing import Any, Dict, List, Optional, Tuple import cv2 import numpy as np @@ -14,6 +14,12 @@ @MODEL_REGISTRY.register(name=ModelType.VSRBaseModel) class VSRBaseModel(CCBaseModel): + def get_bf16_preflight_inputs(self) -> Optional[Tuple[Tuple[Any, ...], Dict[str, Any]]]: + cfg: VSRBaseConfig = self.config + height, width = self._get_bf16_preflight_image_size() + imgs = torch.zeros((1, cfg.num_frame, 3, height, width), device=self.device, dtype=self.half_dtype) + return (imgs,), {} + @torch.inference_mode() # type: ignore def inference(self, img: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: cfg: VSRBaseConfig = self.config @@ -48,8 +54,8 @@ def inference_image_list(self, img_list: List[np.ndarray], *args: Any, **kwargs: # b, n, c, h, w img_tensor_stack = torch.stack(new_img_list, dim=1) - if self.fp16: - img_tensor_stack = img_tensor_stack.half() + if self.fp16 or self.bf16: + img_tensor_stack = img_tensor_stack.to(self.half_dtype) out = self.inference(img_tensor_stack) @@ -58,7 +64,7 @@ def inference_image_list(self, img_list: List[np.ndarray], *args: Any, **kwargs: for i in range(out.shape[1]): img = out[0, i, :, :, :] - img = img.permute(1, 2, 0).cpu().numpy() + img = self._tensor_to_numpy(img.permute(1, 2, 0)) img = (img * 255).clip(0, 255).astype("uint8") img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) res_img_list.append(img) @@ -66,7 +72,7 @@ def inference_image_list(self, img_list: List[np.ndarray], *args: Any, **kwargs: return res_img_list elif len(out.shape) == 4: - img = out.squeeze(0).permute(1, 2, 0).cpu().numpy() + img = self._tensor_to_numpy(out.squeeze(0).permute(1, 2, 0)) img = (img * 255).clip(0, 255).astype("uint8") img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR) diff --git a/tests/sr/test_sr.py b/tests/sr/test_sr.py index fcf836f..bfaa24d 100644 --- a/tests/sr/test_sr.py +++ b/tests/sr/test_sr.py @@ -11,6 +11,7 @@ CCCV_DEVICE, CCCV_FP16, CCCV_TILE, + CI_ENV, calculate_image_similarity, compare_image_size, load_image, @@ -46,6 +47,22 @@ def test_sr_fp16() -> None: assert compare_image_size(img1, img2, cfg.scale) +@pytest.mark.skipif(not CI_ENV, reason="Skip bf16 inference test outside CI") +def test_sr_bf16() -> None: + img1 = load_image() + k = ConfigType.RealESRGAN_AnimeJaNai_HD_V3_Compact_2x + + cfg: BaseConfig = AutoConfig.from_pretrained(k) + model: SRBaseModel = AutoModel.from_config(config=cfg, device=CCCV_DEVICE, fp16=False, bf16=True, tile=CCCV_TILE) + + img2 = model.inference_image(img1) + + cv2.imwrite(str(ASSETS_PATH / f"test_bf16_{k}_out.jpg"), img2) + + assert calculate_image_similarity(img1, img2) + assert compare_image_size(img1, img2, cfg.scale) + + @pytest.mark.skipif( sys.platform == "win32" or not torch_2_4, reason="Skip test torch.compile on Windows or PyTorch version is not 2.4" ) diff --git a/tests/sr/test_transformer_arch.py b/tests/sr/test_transformer_arch.py new file mode 100644 index 0000000..c1829ec --- /dev/null +++ b/tests/sr/test_transformer_arch.py @@ -0,0 +1,135 @@ +from typing import Callable + +import pytest +import torch +from torch import nn + +from cccv.arch.sr.dat_arch import DAT, Spatial_Attention +from cccv.arch.sr.hat_arch import HAT, WindowAttention as HATWindowAttention +from cccv.arch.sr.swinir_arch import SwinIR, WindowAttention as SwinIRWindowAttention + + +def _relative_position_index(window_size: int) -> torch.Tensor: + coords_h = torch.arange(window_size) + coords_w = torch.arange(window_size) + coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij")) + coords_flatten = torch.flatten(coords, 1) + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] + relative_coords = relative_coords.permute(1, 2, 0).contiguous() + relative_coords[:, :, 0] += window_size - 1 + relative_coords[:, :, 1] += window_size - 1 + relative_coords[:, :, 0] *= 2 * window_size - 1 + return relative_coords.sum(-1) + + +def _assert_bf16_attention_forward(forward: Callable[[], torch.Tensor]) -> None: + try: + out = forward() + except RuntimeError as e: + message = str(e) + lower_message = message.lower() + if "bfloat16" in lower_message and ("not implemented" in lower_message or "unsupported" in lower_message): + pytest.skip(f"bfloat16 attention is not supported by this torch build: {message}") + raise + + assert out.dtype == torch.bfloat16 + assert torch.isfinite(out.float()).all() + + +def test_dat_attention_mask_matches_bf16_dtype() -> None: + def forward() -> torch.Tensor: + attn = Spatial_Attention(dim=8, idx=0, split_size=[4, 4], num_heads=2).to(torch.bfloat16).eval() + qkv = torch.randn((3, 1, 16, 8), dtype=torch.bfloat16) + mask = torch.zeros((1, 16, 16), dtype=torch.float32) + with torch.inference_mode(): + return attn(qkv, 4, 4, mask=mask) + + _assert_bf16_attention_forward(forward) + + +def test_dat_relative_position_bias_matches_bf16_dtype() -> None: + def forward() -> torch.Tensor: + attn = Spatial_Attention(dim=8, idx=0, split_size=[4, 4], num_heads=2).eval() + qkv = torch.randn((3, 1, 16, 8), dtype=torch.bfloat16) + with torch.inference_mode(): + return attn(qkv, 4, 4) + + _assert_bf16_attention_forward(forward) + + +def test_hat_attention_mask_matches_bf16_dtype() -> None: + def forward() -> torch.Tensor: + attn = HATWindowAttention(dim=8, window_size=(4, 4), num_heads=2).to(torch.bfloat16).eval() + x = torch.randn((1, 16, 8), dtype=torch.bfloat16) + mask = torch.zeros((1, 16, 16), dtype=torch.float32) + with torch.inference_mode(): + return attn(x, rpi=_relative_position_index(4), mask=mask) + + _assert_bf16_attention_forward(forward) + + +def test_swinir_attention_mask_matches_bf16_dtype() -> None: + def forward() -> torch.Tensor: + attn = SwinIRWindowAttention(dim=8, window_size=(4, 4), num_heads=2).to(torch.bfloat16).eval() + x = torch.randn((1, 16, 8), dtype=torch.bfloat16) + mask = torch.zeros((1, 16, 16), dtype=torch.float32) + with torch.inference_mode(): + return attn(x, mask=mask) + + _assert_bf16_attention_forward(forward) + + +def _dat_model() -> nn.Module: + return DAT( + img_size=16, + embed_dim=16, + split_size=[4, 8], + depth=[3], + num_heads=[4], + expansion_factor=2, + drop_path_rate=0.0, + scale=2, + upsampler="pixelshuffledirect", + ).eval() + + +def _hat_model() -> nn.Module: + return HAT( + img_size=16, + embed_dim=16, + depth=[2], + num_heads=[4], + window_size=4, + compress_ratio=4, + squeeze_factor=4, + mlp_ratio=2, + drop_path_rate=0.0, + scale=2, + upsampler="pixelshuffle", + ).eval() + + +def _swinir_model() -> nn.Module: + return SwinIR( + img_size=16, + embed_dim=16, + depths=[2], + num_heads=[4], + window_size=4, + mlp_ratio=2, + drop_path_rate=0.0, + scale=2, + upsampler="pixelshuffledirect", + ).eval() + + +@pytest.mark.parametrize("make_model", [_dat_model, _hat_model, _swinir_model]) +def test_transformer_arch_forward_shape(make_model: Callable[[], nn.Module]) -> None: + model = make_model() + x = torch.rand((1, 3, 16, 16)) + + with torch.inference_mode(): + out = model(x) + + assert out.shape == (1, 3, 32, 32) + assert out.dtype == torch.float32 diff --git a/tests/test_bf16_fallback.py b/tests/test_bf16_fallback.py new file mode 100644 index 0000000..a1a0a2a --- /dev/null +++ b/tests/test_bf16_fallback.py @@ -0,0 +1,94 @@ +from typing import Any, Dict, Tuple + +import numpy as np +import torch +from torch import nn + +from cccv.model.base_model import CCBaseModel + + +class _Bf16FailingModule(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(1)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if x.dtype == torch.bfloat16: + raise RuntimeError("bf16 kernel unavailable") + return x + self.weight.to(x.dtype) + + +class _Bf16CastFailingModule(nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(1)) + + def to(self, *args: Any, **kwargs: Any) -> "_Bf16CastFailingModule": + dtype = kwargs.get("dtype") + if dtype is None: + dtype = next((arg for arg in args if isinstance(arg, torch.dtype)), None) + if dtype == torch.bfloat16: + raise RuntimeError("bf16 cast unavailable") + return super().to(*args, **kwargs) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self.weight.to(x.dtype) + + +class _Bf16PreflightModel(CCBaseModel): + def __init__(self, *args: Any, **kwargs: Any) -> None: + self.load_count = 0 + super().__init__(*args, **kwargs) + + def load_model(self) -> nn.Module: + self.load_count += 1 + return _Bf16FailingModule().eval().to(self.device) + + def get_bf16_preflight_inputs(self) -> Tuple[Tuple[Any, ...], Dict[str, Any]]: + img = torch.zeros((1, 1), device=self.device, dtype=self.half_dtype) + return (img,), {} + + def inference(self, img: torch.Tensor, *args: object, **kwargs: object) -> torch.Tensor: + return self.model(img) + + +class _Bf16CastFallbackModel(_Bf16PreflightModel): + def load_model(self) -> nn.Module: + self.load_count += 1 + return _Bf16CastFailingModule().eval().to(self.device) + + +def test_bf16_cast_falls_back_to_fp16_when_requested() -> None: + model = _Bf16CastFallbackModel(config=object(), device=torch.device("cpu"), fp16=True, bf16=True) + + assert model.load_count == 2 + assert model.bf16 is False + assert model.fp16 is True + assert model.half_dtype == torch.float16 + assert next(model.model.parameters()).dtype == torch.float16 + + +def test_bf16_preflight_falls_back_to_fp32_when_fp16_disabled() -> None: + model = _Bf16PreflightModel(config=object(), device=torch.device("cpu"), fp16=False, bf16=True) + + assert model.load_count == 2 + assert model.bf16 is False + assert model.fp16 is False + assert model.half_dtype == torch.float32 + assert next(model.model.parameters()).dtype == torch.float32 + + +def test_bf16_preflight_falls_back_to_fp16_when_requested() -> None: + model = _Bf16PreflightModel(config=object(), device=torch.device("cpu"), fp16=True, bf16=True) + + assert model.load_count == 2 + assert model.bf16 is False + assert model.fp16 is True + assert model.half_dtype == torch.float16 + assert next(model.model.parameters()).dtype == torch.float16 + + +def test_tensor_to_numpy_converts_only_bf16_to_float32() -> None: + assert CCBaseModel._tensor_to_numpy(torch.ones((1,), dtype=torch.float32)).dtype == np.float32 + assert CCBaseModel._tensor_to_numpy(torch.ones((1,), dtype=torch.float16)).dtype == np.float16 + assert CCBaseModel._tensor_to_numpy(torch.ones((1,), dtype=torch.bfloat16)).dtype == np.float32 diff --git a/tests/test_util.py b/tests/test_util.py index fa9ff81..c2183c8 100644 --- a/tests/test_util.py +++ b/tests/test_util.py @@ -95,11 +95,13 @@ def test_create_window_3d() -> None: def test_ssim_matlab() -> None: - img1 = torch.randn(1, 3, 64, 64) - img2 = torch.randn(1, 3, 64, 64) - ssim_value = ssim_matlab(img1, img2) - assert isinstance(ssim_value, torch.Tensor) - assert 0.0 <= ssim_value.item() <= 1.0 + img1 = torch.rand(1, 3, 64, 64) + identical_ssim = ssim_matlab(img1, img1.clone()) + similar_ssim = ssim_matlab(img1, (img1 * 0.9 + 0.05).clamp(0, 1)) + + assert isinstance(identical_ssim, torch.Tensor) + assert bool(torch.isclose(identical_ssim, torch.tensor(1.0), atol=1e-4).item()) + assert 0.0 <= similar_ssim.item() <= 1.0 class Test_Check_Scene: