From b904b89755ee619d78a76bb56cea35317cd964cd Mon Sep 17 00:00:00 2001 From: NULL204 Date: Wed, 17 Jun 2026 00:44:36 +0800 Subject: [PATCH] feat: add precision option (fp32/fp16/bf16) Add a precision config field (fp32/fp16/bf16, default fp32) and map it to cccv's fp16/bf16 arguments. bf16 avoids the fp16 numerical overflow that yields NaN output on some transformer models, at the same VRAM savings as fp16. Co-Authored-By: Claude Opus 4.8 (1M context) --- Final2x_core/SRclass.py | 3 ++- Final2x_core/config.py | 8 ++++++++ scripts/gen_config.py | 1 + tests/test_config.py | 25 +++++++++++++++++++++++++ 4 files changed, 36 insertions(+), 1 deletion(-) diff --git a/Final2x_core/SRclass.py b/Final2x_core/SRclass.py index d483aab..a830afb 100644 --- a/Final2x_core/SRclass.py +++ b/Final2x_core/SRclass.py @@ -26,7 +26,8 @@ def __init__(self, config: SRConfig) -> None: self._SR_class: SRBaseModel = AutoModel.from_pretrained( self.config.pretrained_model_name, device=get_device(self.config.device), - fp16=False, + fp16=(self.config.precision == "fp16"), + bf16=(self.config.precision == "bf16"), tile=tile, gh_proxy=self.config.gh_proxy, ) diff --git a/Final2x_core/config.py b/Final2x_core/config.py index 10f592d..20044a2 100644 --- a/Final2x_core/config.py +++ b/Final2x_core/config.py @@ -12,6 +12,7 @@ class SRConfig(BaseModel): pretrained_model_name: Union[ConfigType, str] device: str use_tile: Optional[bool] = None + precision: str = "fp32" # "fp32" | "fp16" | "bf16" gh_proxy: Optional[str] = None target_scale: Optional[Union[int, float]] = None output_path: DirectoryPath @@ -61,3 +62,10 @@ def device_match(cls, v: str) -> str: return v raise ValueError(f"device must start with {device_list}") + + @field_validator("precision") + def precision_match(cls, v: str) -> str: + precision_list = ["fp32", "fp16", "bf16"] + if v not in precision_list: + raise ValueError(f"precision must be one of {precision_list}") + return v diff --git a/scripts/gen_config.py b/scripts/gen_config.py index 15e83b2..798daa4 100644 --- a/scripts/gen_config.py +++ b/scripts/gen_config.py @@ -22,6 +22,7 @@ def gen_config() -> None: "pretrained_model_name": ConfigType.RealESRGAN_AnimeJaNai_HD_V3_Compact_2x.value, "device": _DEVICE_, "use_tile": True, + "precision": "fp32", "gh_proxy": None, "target_scale": None, "output_path": str(projectPATH / "assets"), diff --git a/tests/test_config.py b/tests/test_config.py index a54fbd1..528fc33 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -26,6 +26,31 @@ def test_from_base64(self) -> None: config = SRConfig.from_base64(b_str) print(config) + def test_precision_fp16(self) -> None: + config: SRConfig = SRConfig.from_yaml(CONFIG_PATH) + config.precision = "fp16" + config = SRConfig.from_json_str(config.model_dump_json()) + assert config.precision == "fp16" + + def test_precision_bf16(self) -> None: + config: SRConfig = SRConfig.from_yaml(CONFIG_PATH) + config.precision = "bf16" + config = SRConfig.from_json_str(config.model_dump_json()) + assert config.precision == "bf16" + + def test_precision_default(self) -> None: + config: SRConfig = SRConfig.from_yaml(CONFIG_PATH) + config_dict = config.model_dump() + config_dict.pop("precision", None) + config = SRConfig(**config_dict) + assert config.precision == "fp32" + + def test_precision_invalid(self) -> None: + with pytest.raises(ValueError): + config: SRConfig = SRConfig.from_yaml(CONFIG_PATH) + config.precision = "invalid" + SRConfig.from_json_str(config.model_dump_json()) + def test_error_device(self) -> None: config: SRConfig with pytest.raises(ValueError):