Add support for ESRGAN and custom upscale ratios

scripts/convert_esrgan_pth.py

@@ -0,0 +1,67 @@
+"""Convert a Real-ESRGAN ``*.pth`` checkpoint to ``.safetensors`` for thenoise.
+
+The official ``xinntao/Real-ESRGAN`` releases ship ``*.pth`` torch pickles whose
+state dict is wrapped under a ``params_ema`` (or ``params``) key, which
+thenoise's safetensors-only loader cannot open. This unwraps one and saves it as
+a ``.safetensors`` file that ``thenoise.upscale.esrgan`` accepts directly.
+
+The official weights already use the ComfyUI key naming (``body.*``,
+``conv_first``, ``conv_body``, ``conv_up1``, ``conv_up2``, ``conv_hr``,
+``conv_last``) that the loader expects, so no key remapping is needed — only the
+wrapper is removed.
+
+Usage:
+    python scripts/convert_esrgan_pth.py RealESRGAN_x2plus.pth -o models/esrgan
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors.torch import save_file
+
+# Wrapper keys used by the official releases, in preference order.
+_WRAPPERS = ("params_ema", "params", "state_dict")
+
+
+def _unwrap(obj: Any) -> dict:
+    """Return the bare state dict from a torch-loaded checkpoint ``obj``."""
+    if not isinstance(obj, dict):
+        return obj
+    for key in _WRAPPERS:
+        if key in obj:
+            return obj[key]
+    return obj
+
+
+def convert(pth: Path, out: Path) -> Path:
+    """Load ``pth`` (torch pickle), unwrap it and save as safetensors at ``out``."""
+    state = _unwrap(torch.load(pth, map_location="cpu", weights_only=False))
+    save_file(state, str(out))
+    return out
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(
+        description="Convert a Real-ESRGAN .pth checkpoint to .safetensors"
+    )
+    ap.add_argument("pth", type=Path, help="source .pth checkpoint")
+    ap.add_argument(
+        "-o",
+        "--out",
+        type=Path,
+        default=None,
+        help="output .safetensors path (default: same name next to the .pth)",
+    )
+    args = ap.parse_args()
+
+    out = args.out or args.pth.with_suffix(".safetensors")
+    out.parent.mkdir(parents=True, exist_ok=True)
+    convert(args.pth, out)
+    print(f"saved ESRGAN safetensors to {out}")
+
+
+if __name__ == "__main__":
+    main()

scripts/download_esrgan.py

@@ -0,0 +1,39 @@
+"""Download the Real-ESRGAN x4 model (ComfyUI repackage) for pixel upscaling.
+
+Source: https://huggingface.co/Comfy-Org/Real-ESRGAN_repackaged
+  RealESRGAN_x4plus.safetensors
+
+Used by thenoise's ``fast`` upscale path, and by the ``refined`` path when
+``--upscale-factor`` exceeds the latent 2x. Optional: if the model is absent
+only the refiner (latent) upscale is available.
+
+Usage:
+    python scripts/download_esrgan.py --out ./models/esrgan
+"""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from huggingface_hub import hf_hub_download
+
+REPO = "Comfy-Org/Real-ESRGAN_repackaged"
+FILE = "RealESRGAN_x4plus.safetensors"
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description="Download the Real-ESRGAN x4 model")
+    ap.add_argument(
+        "--out", default="./models/esrgan", help="output directory"
+    )
+    args = ap.parse_args()
+
+    out = Path(args.out)
+    out.mkdir(parents=True, exist_ok=True)
+
+    path = hf_hub_download(repo_id=REPO, filename=FILE, local_dir=out)
+    print(f"saved ESRGAN model to {path}")
+
+
+if __name__ == "__main__":
+    main()

tests/test_upscale.py

@@ -0,0 +1,108 @@
+"""Upscale factor/type planning tests (no torch, no weights)."""
+from __future__ import annotations
+
+import pytest
+
+from thenoise.models.base import DiffusionModel
+
+
+def _make_model(esrgan_path=None, esrgan_scale=None):
+    """Build a minimal concrete subclass (bypassing __init__ / VAE)."""
+
+    class _M(DiffusionModel):
+        name = "test"
+
+        @staticmethod
+        def detect(f):
+            return False
+
+        def encode_prompt(self, prompt, negative_prompt="", *, guidance_scale):
+            pass
+
+        def init_latents(self, height, width, seed):
+            pass
+
+        def schedule(self, steps, height, width):
+            pass
+
+        def denoise_step(self, latents, t, cond, guidance_scale, i):
+            pass
+
+        def _upscale_format(self):
+            return "wan21"
+
+    m = object.__new__(_M)
+    m.esrgan_path = esrgan_path
+    m._esrgan_scale_val = esrgan_scale
+    return m
+
+
+def test_esrgan_scale_mapping_with_4x():
+    m = _make_model(esrgan_path="/tmp/x4.safetensors", esrgan_scale=4)
+    # refined: latent gives 2x, ESRGAN 4x only above the latent 2x.
+    assert m._esrgan_scale_for(1.0, "refined") == 0
+    assert m._esrgan_scale_for(2.0, "refined") == 0
+    assert m._esrgan_scale_for(2.5, "refined") == 4
+    assert m._esrgan_scale_for(8.0, "refined") == 4
+    assert m._esrgan_scale_for(0.5, "refined") == 0
+    # fast: no latent multiplier, always ESRGAN for any upscale.
+    assert m._esrgan_scale_for(1.5, "fast") == 4
+    assert m._esrgan_scale_for(4.0, "fast") == 4
+    assert m._esrgan_scale_for(0.5, "fast") == 0
+
+
+def test_esrgan_scale_mapping_with_2x():
+    m = _make_model(esrgan_path="/tmp/x2.safetensors", esrgan_scale=2)
+    assert m._esrgan_scale_for(2.0, "refined") == 0
+    assert m._esrgan_scale_for(2.5, "refined") == 2
+    assert m._esrgan_scale_for(4.0, "refined") == 2
+    assert m._esrgan_scale_for(1.5, "fast") == 2
+    assert m._esrgan_scale_for(2.0, "fast") == 2
+
+
+def test_resolve_valid_refined_without_esrgan():
+    m = _make_model(esrgan_path=None)
+    # f <= latent 2x in refined mode needs no ESRGAN.
+    assert m._resolve_upscale(1.0, "refined") == (1.0, "refined")
+    assert m._resolve_upscale(2.0, "refined") == (2.0, "refined")
+
+
+def test_resolve_needs_esrgan_when_absent():
+    m = _make_model(esrgan_path=None)
+    with pytest.raises(ValueError):
+        m._resolve_upscale(2.5, "refined")
+    with pytest.raises(ValueError):
+        m._resolve_upscale(1.5, "fast")
+
+
+def test_resolve_max_ranges_depend_on_scale():
+    # 4x model: refined up to 8, fast up to 4.
+    m4 = _make_model("/tmp/x4.safetensors", esrgan_scale=4)
+    assert m4._resolve_upscale(8.0, "refined") == (8.0, "refined")
+    assert m4._resolve_upscale(4.0, "fast") == (4.0, "fast")
+    with pytest.raises(ValueError):
+        m4._resolve_upscale(5.0, "fast")
+    with pytest.raises(ValueError):
+        m4._resolve_upscale(9.0, "refined")
+
+    # 2x model: refined up to 4, fast up to 2.
+    m2 = _make_model("/tmp/x2.safetensors", esrgan_scale=2)
+    assert m2._resolve_upscale(4.0, "refined") == (4.0, "refined")
+    assert m2._resolve_upscale(2.0, "fast") == (2.0, "fast")
+    with pytest.raises(ValueError):
+        m2._resolve_upscale(5.0, "refined")
+    with pytest.raises(ValueError):
+        m2._resolve_upscale(3.0, "fast")
+
+
+def test_resolve_invalid_factor():
+    m = _make_model("/tmp/x4.safetensors", esrgan_scale=4)
+    for bad in (0.0, -1.0, 8.5):
+        with pytest.raises(ValueError):
+            m._resolve_upscale(bad, "refined")
+
+
+def test_resolve_invalid_type():
+    m = _make_model("/tmp/x4.safetensors", esrgan_scale=4)
+    with pytest.raises(ValueError):
+        m._resolve_upscale(2.0, "bogus")

thenoise/__main__.py

@@ -15,6 +15,7 @@ def _serve(args) -> None:
             vae_path=args.vae,
             text_encoder_path=args.text_encoder,
             lora_dir=args.lora_dir,
+            esrgan_path=args.esrgan,
         ),
     )
 

thenoise/api.py

@@ -38,6 +38,8 @@ class Text2ImageRequest(BaseModel):
     guidance_scale: Optional[float] = None
     seed: Optional[int] = None
     upscale: bool = False
+    upscale_factor: float = 1.0
+    upscale_type: str = "refined"
     sampler: Optional[str] = None
     qwen_vae_enhance: bool = False
     film_grain: float = 0.0
@@ -83,6 +85,8 @@ def text2image(req: Text2ImageRequest):
                 guidance_scale=req.guidance_scale,
                 seed=req.seed,
                 upscale=req.upscale,
+                upscale_factor=req.upscale_factor,
+                upscale_type=req.upscale_type,
                 sampler=req.sampler,
                 qwen_vae_enhance=req.qwen_vae_enhance,
                 film_grain=req.film_grain,

thenoise/cli.py

@@ -23,6 +23,10 @@ def _add_model_paths(p: argparse.ArgumentParser) -> None:
     p.add_argument("--lora-dir", default="", metavar="PATH",
                    help="directory containing LoRA .safetensors files "
                         "(subdirectories allowed)")
+    p.add_argument("--esrgan", default="", metavar="PATH",
+                   help="optional Real-ESRGAN model (.safetensors) for pixel "
+                        "upscaling; without it only refiner (latent) upscale is "
+                        "available")
 
 
 def build_parser() -> argparse.ArgumentParser:
@@ -57,6 +61,15 @@ def build_parser() -> argparse.ArgumentParser:
     gen.add_argument("--upscale", action="store_true",
                      help="upscale the latent 2x in latent space (SesquiLSR) and "
                           "run a low-strength refine denoise before decoding")
+    gen.add_argument("--upscale-factor", type=float, default=1.0,
+                     help="upscale factor, > 0.0 (default: 1.0 = no upscale); "
+                          "max depends on the ESRGAN model scale: 'fast' is "
+                          "limited to the model scale, 'refined' to latent 2x * "
+                          "model scale")
+    gen.add_argument("--upscale-type", choices=["refined", "fast"],
+                     default="refined",
+                     help="'refined' (default): latent 2x + refiner, plus ESRGAN "
+                          "above factor 2; 'fast': ESRGAN only (no latent 2x)")
     gen.add_argument("--sampler", choices=["euler", "er_sde"], default=None,
                      help="denoising solver (default: er_sde)")
     gen.add_argument("--qwen-vae-enhance", action="store_true",

thenoise/generate.py

@@ -24,6 +24,7 @@ def run_generate(args) -> None:
             vae_path=args.vae,
             text_encoder_path=args.text_encoder,
             lora_dir=args.lora_dir,
+            esrgan_path=args.esrgan,
         ),
     )
 
@@ -37,6 +38,8 @@ def run_generate(args) -> None:
         guidance_scale=args.guidance_scale,
         seed=seed,
         upscale=args.upscale,
+        upscale_factor=args.upscale_factor,
+        upscale_type=args.upscale_type,
         sampler=args.sampler,
         qwen_vae_enhance=args.qwen_vae_enhance,
         film_grain=args.film_grain,

thenoise/models/anima.py

@@ -58,6 +58,7 @@ def __init__(
         device: str = "cuda",
         dtype: torch.dtype = torch.bfloat16,
         lora_dir: Optional[str] = None,
+        esrgan_path: Optional[str] = None,
     ):
         super().__init__(
             dit_path=dit_path,
@@ -66,6 +67,7 @@ def __init__(
             device=device,
             dtype=dtype,
             lora_dir=lora_dir,
+            esrgan_path=esrgan_path,
         )
 
         logger.info("Loading Anima DiT from %s", dit_path)