Upload folder using huggingface_hub

README.md

@@ -31,7 +31,7 @@ the two weights this project publishes:
 | Model | Cosa è / What it is | Quando sceglierlo / Pick it when |
 |---|---|---|
 | 🔍 **Universal Upscale** (default) | LIIF v4 GAN, 4.60 M parametri, `model_scripted.pt` con tiled inference integrata | qualsiasi fattore 1x-9.2x o una dimensione target esatta, immagini grandi |
-| ⚡ **Universal Upscale Distilled** | student knowledge-distilled, 1.30 M parametri, ONNX single-file `model.onnx` (5.2 MB) | serve 9x in un solo passaggio: più leggero e veloce |
+| ⚡ **Universal Upscale Distilled** | student knowledge-distilled, 1.30 M parametri, TorchScript `model_distilled.pt` che gira sullo stesso device del modello di default | serve 9x in un solo passaggio: più leggero e veloce |
 
 Lo Student distillato ha una scala nativa **fissa**: sotto 9x l'output viene
 ricampionato in area-downsample, sopra viene ricampionato in bicubico. La riga di stato indica
@@ -42,7 +42,7 @@ quale modello ha girato e se l'output è stato ricampionato. Model card:
 
 1. **Carica un'immagine / Upload an image**: trascina o carica un'immagine dal tuo dispositivo.
 2. **Scegli il modello / Pick the model**: lascia **Universal Upscale** (default) oppure passa a
-   **Universal Upscale Distilled** (ONNX, nativo 9x).
+   **Universal Upscale Distilled** (nativo 9x, usa la stessa GPU del modello di default).
 3. **Imposta la scala / Set the scale**:
    - **Scale Factor**: scegli un fattore di scala da **1x fino a 9.2x** (passo 0.2x, es. 2x, 3.6x, 4x, 9.2x).
    - **Target Long Side (px)**: in alternativa, specifica la dimensione in pixel del lato più lungo desiderato (es. 2048, 3000, 4096).

app.py

@@ -28,8 +28,11 @@ except ImportError:
         return fn
 
 MODEL_PATH = Path(__file__).resolve().parent / "model_scripted.pt"
-DISTILLED_ONNX_PATH = Path(__file__).resolve().parent / "model.onnx"
-# If the Space folder does not carry model.onnx, it is pulled from the distilled repo on first use.
+# The distilled student ships beside it as its own TorchScript file: the same torch runtime the
+# default engine already uses, so the student lands on the Space GPU. The ONNX build was dropped
+# because the Space installs the CPU-only `onnxruntime` wheel, which pinned it to CPU.
+DISTILLED_PT_PATH = Path(__file__).resolve().parent / "model_distilled.pt"
+# If the Space folder does not carry model_distilled.pt, it is pulled from the distilled repo on first use.
 DISTILLED_REPO_ID = "tailot/universal-upscale-distilled"
 
 MAX_SIDE = 4096
@@ -44,7 +47,7 @@ DISTILLED_NATIVE = 9
 # Plain strings (not (label, value) tuples) so the dropdown works on any Gradio 5 minor version;
 # upscale() maps the label back to a model key.
 MODEL_LIIF = "🔍 Universal Upscale - arbitrary scale 1x-9.2x (default)"
-MODEL_DISTILLED = "⚡ Universal Upscale Distilled - ONNX, native 9x"
+MODEL_DISTILLED = "⚡ Universal Upscale Distilled - fixed 9x, runs on GPU"
 MODEL_CHOICES = [MODEL_LIIF, MODEL_DISTILLED]
 
 
@@ -52,11 +55,11 @@ def normalize_model_choice(value) -> str:
     """Map the dropdown selection to a model key.
 
     Gradio pads short example rows with None, and a Dropdown can hand back either the label or
-    the value, so the mapping is deliberately loose: anything mentioning "distilled" is the ONNX
-    student, everything else stays on the published default.
+    the value, so the mapping is deliberately loose: anything mentioning "distilled" (or the old
+    ONNX labels) is the fixed-factor student, everything else stays on the published default.
     """
     text = str(value or "").strip().lower()
-    return "distilled" if "distill" in text or text in ("onnx", "student") else "liif"
+    return "distilled" if "distill" in text or text in ("onnx", "student", "torchscript") else "liif"
 
 
 def load_model(device: str):
@@ -76,61 +79,64 @@ def resolve_factor(input_size, scale, target_long_side):
     return factor, ""
 
 
-_ONNX_SESSION = None
-_ONNX_NATIVE = 0
+_DISTILLED = None  # (model, native factor, device), built once per container
 
 
-def onnx_session():
-    """ONNX Runtime session for the distilled student, built once per container.
+def distilled_model():
+    """(model, native factor, device) for the distilled student, built once per container.
 
-    The memory arena and the memory pattern are switched off on purpose: tiled inference hands
-    ONNX Runtime a different shape at every tile, and its pattern allocator keeps the biggest
-    tile's buffers alive until the Space runs out of RAM.
+    The student is loaded with the same torch runtime as the default engine, so it runs on the
+    Space GPU whenever there is one. ONNX Runtime was dropped on purpose: the Space installs the
+    CPU-only `onnxruntime` wheel, so the ONNX build of these weights could never reach the GPU.
     """
-    global _ONNX_SESSION, _ONNX_NATIVE
-    if _ONNX_SESSION is not None:
-        return _ONNX_SESSION
-    import onnxruntime as ort
-
-    path = DISTILLED_ONNX_PATH
+    global _DISTILLED
+    if _DISTILLED is not None:
+        return _DISTILLED
+    path = DISTILLED_PT_PATH
     if not path.exists():
         from huggingface_hub import hf_hub_download
 
-        path = Path(hf_hub_download(repo_id=DISTILLED_REPO_ID, filename="model.onnx"))
-    options = ort.SessionOptions()
-    options.enable_cpu_mem_arena = False
-    options.enable_mem_pattern = False
-    options.log_severity_level = 3
-    available = ort.get_available_providers()
-    providers = [p for p in ("CUDAExecutionProvider", "CPUExecutionProvider") if p in available]
-    session = ort.InferenceSession(str(path), sess_options=options, providers=providers or None)
-    # Read the fixed factor from the model itself instead of trusting a constant.
-    probe = session.run(None, {session.get_inputs()[0].name: np.zeros((1, 3, 8, 9), dtype=np.float32)})[0]
-    _ONNX_NATIVE, _ONNX_SESSION = int(round(max(probe.shape[-2] / 8.0, probe.shape[-1] / 9.0))), session
-    return session
-
-
-def upscale_distilled(rgb, scale, native, out_w, out_h, tile: int = 64, pad: int = 6):
+        path = Path(hf_hub_download(repo_id=DISTILLED_REPO_ID, filename="model_scripted.pt"))
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = torch.jit.load(str(path), map_location=device)
+    model.eval()
+    with torch.no_grad():
+        # Read the fixed factor from the model itself instead of trusting a constant.
+        probe = model(torch.zeros(1, 3, 8, 9, device=device))
+        native = int(round(max(probe.shape[-2] / 8.0, probe.shape[-1] / 9.0)))
+        # Warm the kernels with one full window (tile 64 + 2*6 padding, the defaults of
+        # upscale_distilled below) so the first real request does not pay for cuDNN kernel
+        # selection, which costs seconds on a cold GPU container.
+        model(torch.zeros(1, 3, 76, 76, device=device))
+    _DISTILLED = (model, native, device)
+    return _DISTILLED
+
+
+def upscale_distilled(rgb, scale, native, out_w, out_h, model=None, device="cpu",
+                      tile: int = 64, pad: int = 6):
     """Tiled inference in LR space for the fixed-factor student.
 
     Same recipe as scripts/upscale_onnx.py: run the padded tile, resample the WHOLE padded tile
     to its own target size (so the crop offset stays inside it), then paste the crop on the canvas.
     """
-    session = onnx_session()
-    input_name = session.get_inputs()[0].name
     height, width = rgb.shape[:2]
     canvas = np.zeros((out_h, out_w, 3), dtype=np.uint8)
     resample = Image.Resampling.LANCZOS if native > scale else Image.Resampling.BICUBIC
     tile = max(8, min(int(tile), height))
+    # Every window keeps the SAME size: at the borders it slides inward instead of shrinking.
+    # TorchScript re-specialises its fused convolutions for the first window of each new shape
+    # (~0.3 s each), so a variable window would make the first request pay for up to nine shapes.
+    win_h, win_w = min(tile + 2 * pad, height), min(tile + 2 * pad, width)
     for top in range(0, height, tile):
         for left in range(0, width, tile):
             bottom, right = min(top + tile, height), min(left + tile, width)
-            y0, y1 = max(top - pad, 0), min(bottom + pad, height)
-            x0, x1 = max(left - pad, 0), min(right + pad, width)
+            y0, x0 = max(0, min(top - pad, height - win_h)), max(0, min(left - pad, width - win_w))
+            y1, x1 = y0 + win_h, x0 + win_w
             patch = rgb[y0:y1, x0:x1].astype(np.float32) / 255.0
-            tensor = np.ascontiguousarray(patch.transpose(2, 0, 1)[None, :, :, :])
-            sr = np.clip(session.run(None, {input_name: tensor})[0][0], 0.0, 1.0)
-            sr = (sr.transpose(1, 2, 0) * 255.0).round().astype(np.uint8)
+            tensor = torch.from_numpy(patch).permute(2, 0, 1).unsqueeze(0).to(device)
+            with torch.no_grad():
+                sr = model(tensor)[0].clamp(0.0, 1.0).permute(1, 2, 0).cpu().numpy()
+            sr = (sr * 255.0).round().astype(np.uint8)
             full_w, full_h = int(round((y1 - y0) * scale)), int(round((x1 - x0) * scale))
             sr = np.asarray(Image.fromarray(sr).resize((full_w, full_h), resample))
             th, tw = int(round((bottom - top) * scale)), int(round((right - left) * scale))
@@ -162,15 +168,15 @@ def upscale(image, model_choice=None, scale: float = 4.0, target_long_side: floa
 
     if choice == "distilled":
         try:
-            session = onnx_session()
+            model, native, device = distilled_model()
         except Exception as exc:
-            return None, f"Distilled ONNX unavailable: {exc}"
-        native = _ONNX_NATIVE or DISTILLED_NATIVE
-        arr = upscale_distilled(np.asarray(img, dtype=np.uint8), applied, native, out_w, out_h)
-        provider = (session.get_providers()[0] or "CPU").replace("ExecutionProvider", "")
-        model_note = f"Universal Upscale Distilled (ONNX, native {native}x"
+            return None, f"Distilled model unavailable: {exc}"
+        native = native or DISTILLED_NATIVE
+        arr = upscale_distilled(np.asarray(img, dtype=np.uint8), applied, native, out_w, out_h,
+                                model=model, device=device)
+        model_note = f"Universal Upscale Distilled (TorchScript, native {native}x"
         model_note += ", output resampled)" if abs(applied - native) > 1e-3 else ")"
-        device = f"onnxruntime/{provider}"
+        device = f"torch/{device}"
     else:
         device = "cuda" if torch.cuda.is_available() else "cpu"
         model = load_model(device)
@@ -209,8 +215,9 @@ demo = gr.Interface(
     description=(
         "Pick the engine in the **Model** menu: **Universal Upscale** renders any factor from 1x up to the tested "
         "maximum 9.2x (or an exact target resolution), **Universal Upscale Distilled** is the 1.3 M "
-        "parameter ONNX student that renders 9x natively in one pass and resamples it for "
-        "other factors. Tiled inference keeps high-resolution upscaling inside memory (outputs are capped at "
+        "parameter student that renders 9x natively in one pass on the same device as the "
+        "default engine and resamples its output for other factors. Tiled inference keeps high-resolution "
+        "upscaling inside memory (outputs are capped at "
         "4096 px on the long side); the status line reports which model ran and what it did."
     ),
     examples=example_files if example_files else None,

config.json

@@ -65,7 +65,7 @@
   },
   "export": {
     "torch_version": "2.14.1+cu130",
-    "exported_at": "2026-10-25T21:52:14+00:00",
+    "exported_at": "2026-10-25T23:09:39+00:00",
     "method": "torch.jit.script + torch.jit.freeze",
     "verification": {
       "checks": [
@@ -169,8 +169,8 @@
       "scales": "arbitrary"
     },
     "distilled": {
-      "file": "model.onnx",
-      "type": "onnx",
+      "file": "model_distilled.pt",
+      "type": "torchscript",
       "repo_id": "tailot/universal-upscale-distilled",
       "native_scale": 9,
       "note": "fixed-factor student; other factors are resampled from the native output"

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3479f51ce4243f71f5c6ffd2d8dc46df82eb39537936df37f30867f3a196e45e
+oid sha256:3d12afef19d94cf6922b8ede382c85db09cd696fbc28ed8ef964b7278343c2d8
 size 18439020

model_distilled.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fbfc0d38849bf09d43cf5f56da50bba63de099ea7bfaff4169741cf90328785f
+size 5292159

requirements.txt

@@ -3,5 +3,4 @@ numpy>=1.26
 Pillow>=10.2
 huggingface_hub>=0.28.1
 gradio>=5.0.0
-onnxruntime>=1.16
 spaces