Add source code

.gitignore

@@ -0,0 +1,35 @@
+# Model weights (large files — download via `python run.py convert` or get from HuggingFace)
+weights/
+
+# Python
+__pycache__/
+*.py[cod]
+*.pyo
+*.pyd
+.Python
+*.egg-info/
+dist/
+build/
+.eggs/
+*.egg
+
+# Virtual environments
+.venv/
+venv/
+env/
+.env
+
+# OS
+.DS_Store
+Thumbs.db
+
+# Editor
+.vscode/
+.idea/
+*.swp
+*.swo
+
+# Output files
+output/
+outputs/
+*.tmp

LICENSE

@@ -0,0 +1,159 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship made available under
+      the License, as indicated by a copyright notice that is included in
+      or attached to the work (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean, as submitted to the Licensor for inclusion
+      in the Work by the copyright owner or by an individual or Legal Entity
+      authorized to submit on behalf of the copyright owner. For the purposes
+      of this definition, "submit" means any form of electronic, verbal, or
+      written communication sent to the Licensor or its representatives,
+      including but not limited to communication on electronic mailing lists,
+      source code control systems, and issue tracking systems that are managed
+      by, or on behalf of, the Licensor for the purpose of developing the Work,
+      but excluding communication that is conspicuously marked or designated in
+      writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any Legal Entity on behalf of
+      whom a Contribution has been received by the Licensor and included
+      within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by the combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a cross-claim
+      or counterfault in a lawsuit) alleging that the Work or any
+      Contribution embodied within the Work constitutes direct or contributory
+      patent infringement, then any patent licenses granted to You under
+      this License for that Work shall terminate as of the date such
+      litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or Derivative
+          Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, You must include a readable copy of the
+          attribution notices contained within such NOTICE file, in
+          at least one of the following places: within a NOTICE text
+          file distributed as part of the Derivative Works; within
+          the Source form or documentation, if provided along with the
+          Derivative Works; or, within a display generated by the
+          Derivative Works, if and wherever such third-party notices
+          normally appear. The contents of the NOTICE file are for
+          informational purposes only and do not modify the License.
+          You may add Your own attribution notices within Derivative
+          Works that You distribute, alongside or in addition to the
+          NOTICE text from the Work, provided that such additional
+          attribution notices cannot be construed as modifying the License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or reproducing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or exemplary damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or all other
+      commercial damages or losses), even if such Contributor has been
+      advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may offer only
+      conditions consistent with this License.
+
+   END OF TERMS AND CONDITIONS

convert.py

@@ -0,0 +1,308 @@
+#!/usr/bin/env python3
+"""
+Convert Thera RDN weights (air/pro) from Flax pickle to MLX safetensors format.
+
+Requires: jax, flax, numpy, safetensors, huggingface_hub
+These are only needed for conversion, not for inference.
+"""
+
+import argparse
+import os
+import pickle
+import sys
+
+import numpy as np
+
+
+def conv_weight(kernel):
+    """Flax Conv (H, W, C_in, C_out) → MLX Conv2d (C_out, H, W, C_in)"""
+    return np.transpose(np.asarray(kernel), (3, 0, 1, 2))
+
+
+def dense_weight(kernel):
+    """Flax Dense (in, out) → MLX Linear (out, in)"""
+    return np.transpose(np.asarray(kernel), (1, 0))
+
+
+def layernorm_params(flax_ln, prefix):
+    """Map Flax LayerNorm scale/bias to MLX weight/bias."""
+    w = {}
+    w[f'{prefix}.weight'] = np.asarray(flax_ln['scale'])
+    if 'bias' in flax_ln:
+        w[f'{prefix}.bias'] = np.asarray(flax_ln['bias'])
+    return w
+
+
+def convert_rdn_encoder(enc):
+    """Convert RDN backbone weights."""
+    weights = {}
+
+    # SFE1 (Conv_0) and SFE2 (Conv_1)
+    weights['encoder.sfe1.weight'] = conv_weight(enc['Conv_0']['kernel'])
+    weights['encoder.sfe1.bias'] = np.asarray(enc['Conv_0']['bias'])
+    weights['encoder.sfe2.weight'] = conv_weight(enc['Conv_1']['kernel'])
+    weights['encoder.sfe2.bias'] = np.asarray(enc['Conv_1']['bias'])
+
+    # 16 Residual Dense Blocks
+    for i in range(16):
+        rdb = enc[f'RDB_{i}']
+        # 8 RDB_Conv layers per block
+        for j in range(8):
+            rc = rdb[f'RDB_Conv_{j}']['Conv_0']
+            prefix = f'encoder.rdbs.{i}.convs.{j}.conv'
+            weights[f'{prefix}.weight'] = conv_weight(rc['kernel'])
+            weights[f'{prefix}.bias'] = np.asarray(rc['bias'])
+        # Local fusion 1x1 conv (Conv_0 at RDB level)
+        lf = rdb['Conv_0']
+        prefix = f'encoder.rdbs.{i}.local_fusion'
+        weights[f'{prefix}.weight'] = conv_weight(lf['kernel'])
+        weights[f'{prefix}.bias'] = np.asarray(lf['bias'])
+
+    # Global Feature Fusion (Conv_2 = 1x1, Conv_3 = 3x3)
+    weights['encoder.gff_1x1.weight'] = conv_weight(enc['Conv_2']['kernel'])
+    weights['encoder.gff_1x1.bias'] = np.asarray(enc['Conv_2']['bias'])
+    weights['encoder.gff_3x3.weight'] = conv_weight(enc['Conv_3']['kernel'])
+    weights['encoder.gff_3x3.bias'] = np.asarray(enc['Conv_3']['bias'])
+
+    return weights
+
+
+def convert_swinir_tail(ref):
+    """Convert SwinIR tail (refine) weights for rdn-pro."""
+    weights = {}
+
+    # conv_first: refine/Conv_0
+    weights['refine.conv_first.weight'] = conv_weight(ref['Conv_0']['kernel'])
+    weights['refine.conv_first.bias'] = np.asarray(ref['Conv_0']['bias'])
+
+    # patch_embed_norm: refine/PatchEmbed_0/LayerNorm_0
+    weights.update(layernorm_params(
+        ref['PatchEmbed_0']['LayerNorm_0'], 'refine.patch_embed_norm'))
+
+    # RSTB layers
+    rstb_depths = [7, 6]  # number of SwinTransformerBlocks per RSTB
+    for i, depth in enumerate(rstb_depths):
+        rstb = ref[f'RSTB_{i}']
+        basic = rstb['BasicLayer_0']
+
+        for j in range(depth):
+            stb = basic[f'SwinTransformerBlock_{j}']
+            mlx_prefix = f'refine.layers.{i}.blocks.{j}'
+
+            # LayerNorm_0 → norm1
+            weights.update(layernorm_params(
+                stb['LayerNorm_0'], f'{mlx_prefix}.norm1'))
+
+            # WindowAttention_0
+            wa = stb['WindowAttention_0']
+            # qkv Dense → Linear
+            weights[f'{mlx_prefix}.attn.qkv.weight'] = dense_weight(wa['qkv']['kernel'])
+            weights[f'{mlx_prefix}.attn.qkv.bias'] = np.asarray(wa['qkv']['bias'])
+            # proj Dense → Linear
+            weights[f'{mlx_prefix}.attn.proj.weight'] = dense_weight(wa['proj']['kernel'])
+            weights[f'{mlx_prefix}.attn.proj.bias'] = np.asarray(wa['proj']['bias'])
+            # relative_position_bias_table (no transform needed)
+            weights[f'{mlx_prefix}.attn.relative_position_bias_table'] = \
+                np.asarray(wa['relative_position_bias_table'])
+
+            # LayerNorm_1 → norm2
+            weights.update(layernorm_params(
+                stb['LayerNorm_1'], f'{mlx_prefix}.norm2'))
+
+            # Mlp_0 → mlp
+            mlp = stb['Mlp_0']
+            weights[f'{mlx_prefix}.mlp.fc1.weight'] = dense_weight(mlp['Dense_0']['kernel'])
+            weights[f'{mlx_prefix}.mlp.fc1.bias'] = np.asarray(mlp['Dense_0']['bias'])
+            weights[f'{mlx_prefix}.mlp.fc2.weight'] = dense_weight(mlp['Dense_1']['kernel'])
+            weights[f'{mlx_prefix}.mlp.fc2.bias'] = np.asarray(mlp['Dense_1']['bias'])
+
+        # RSTB conv: RSTB_{i}/Conv_0
+        weights[f'refine.layers.{i}.conv.weight'] = conv_weight(rstb['Conv_0']['kernel'])
+        weights[f'refine.layers.{i}.conv.bias'] = np.asarray(rstb['Conv_0']['bias'])
+
+    # Final norm: refine/LayerNorm_0
+    weights.update(layernorm_params(ref['LayerNorm_0'], 'refine.norm'))
+
+    # conv_after_body: refine/Conv_1
+    weights['refine.conv_after_body.weight'] = conv_weight(ref['Conv_1']['kernel'])
+    weights['refine.conv_after_body.bias'] = np.asarray(ref['Conv_1']['bias'])
+
+    # conv_last: refine/Conv_2
+    weights['refine.conv_last.weight'] = conv_weight(ref['Conv_2']['kernel'])
+    weights['refine.conv_last.bias'] = np.asarray(ref['Conv_2']['bias'])
+
+    return weights
+
+
+def convert_flax_to_mlx(flax_params, size='air'):
+    """Map Flax parameter tree to flat MLX weight dict."""
+    p = flax_params['params']
+    weights = {}
+
+    # --- Global params (no transform) ---
+    weights['k'] = np.asarray(p['k'], dtype=np.float32).reshape(())
+    weights['components'] = np.asarray(p['components'], dtype=np.float32)
+
+    # --- RDN Backbone ---
+    weights.update(convert_rdn_encoder(p['encoder']))
+
+    # --- SwinIR tail (pro only) ---
+    if size == 'pro':
+        weights.update(convert_swinir_tail(p['refine']))
+
+    # --- Hypernetwork output conv ---
+    weights['out_conv.weight'] = conv_weight(p['out_conv']['kernel'])
+    weights['out_conv.bias'] = np.asarray(p['out_conv']['bias'])
+
+    return weights
+
+
+REPO_IDS = {
+    'air': 'prs-eth/thera-rdn-air',
+    'pro': 'prs-eth/thera-rdn-pro',
+}
+
+
+def download_model(size='air', filename="model.pkl", cache_dir=None):
+    """Download model pickle from HuggingFace."""
+    from huggingface_hub import hf_hub_download
+    repo_id = REPO_IDS[size]
+    return hf_hub_download(repo_id=repo_id, filename=filename, cache_dir=cache_dir)
+
+
+def load_pickle_with_jax(path):
+    """Load Flax pickle (requires jax and flax installed)."""
+    try:
+        import jax  # noqa: F401 - needed for array reconstruction
+        import flax  # noqa: F401 - needed for FrozenDict
+    except ImportError:
+        print("Error: jax and flax are required for loading the original weights.")
+        print("Install them with: pip install jax flax")
+        sys.exit(1)
+
+    with open(path, 'rb') as f:
+        checkpoint = pickle.load(f)
+
+    params = checkpoint['model']
+    backbone = checkpoint['backbone']
+    size = checkpoint['size']
+    print(f"Loaded checkpoint: backbone={backbone}, size={size}")
+
+    if backbone != 'rdn':
+        print(f"Warning: this converter is designed for rdn, got {backbone}")
+
+    return params
+
+
+def load_pickle_without_jax(path):
+    """
+    Attempt to load Flax pickle without JAX by mocking the required classes.
+    Falls back to JAX-based loading if this fails.
+    """
+    import types
+
+    class MockFrozenDict(dict):
+        pass
+
+    class MockModule(types.ModuleType):
+        def __getattr__(self, name):
+            return MockModule(name)
+
+    class NumpyUnpickler(pickle.Unpickler):
+        def find_class(self, module, name):
+            if 'frozen_dict' in module and name == 'FrozenDict':
+                return MockFrozenDict
+            if module.startswith('jax') and name == '_reconstruct_array':
+                # JAX arrays are reconstructed from numpy arrays + metadata
+                def reconstruct(*args):
+                    # args typically: (numpy_array, dtype, weak_type)
+                    if len(args) >= 1 and isinstance(args[0], np.ndarray):
+                        return args[0]
+                    return np.array(args[0])
+                return reconstruct
+            if module.startswith('jax'):
+                try:
+                    return super().find_class(module, name)
+                except (ImportError, AttributeError):
+                    return lambda *a, **kw: a[0] if a else None
+            return super().find_class(module, name)
+
+    try:
+        with open(path, 'rb') as f:
+            checkpoint = NumpyUnpickler(f).load()
+        params = checkpoint['model']
+        backbone = checkpoint['backbone']
+        size = checkpoint['size']
+        print(f"Loaded checkpoint (no-jax mode): backbone={backbone}, size={size}")
+        return params
+    except Exception as e:
+        print(f"Mock unpickle failed ({e}), falling back to JAX-based loading...")
+        return load_pickle_with_jax(path)
+
+
+def save_safetensors(weights, output_path):
+    """Save weight dict as safetensors."""
+    from safetensors.numpy import save_file
+    save_file(weights, output_path)
+    print(f"Saved MLX weights to {output_path}")
+
+
+def save_npz(weights, output_path):
+    """Save weight dict as npz (fallback if safetensors not available)."""
+    np.savez(output_path, **weights)
+    print(f"Saved MLX weights to {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Convert Thera RDN weights to MLX format")
+    parser.add_argument('--model', type=str, choices=['air', 'pro'], default='air',
+                        help='Model variant (default: air)')
+    parser.add_argument('--input', type=str, default=None,
+                        help='Path to model.pkl (downloads from HuggingFace if not provided)')
+    parser.add_argument('--output', type=str, default=None,
+                        help='Output path (default: weights-{model}.safetensors)')
+    parser.add_argument('--no-jax', action='store_true',
+                        help='Try to load pickle without JAX installed')
+    args = parser.parse_args()
+
+    if args.output is None:
+        args.output = f'weights-{args.model}.safetensors'
+
+    # Download if needed
+    if args.input is None:
+        repo = REPO_IDS[args.model]
+        print(f"Downloading model from HuggingFace ({repo})...")
+        pkl_path = download_model(args.model)
+    else:
+        pkl_path = args.input
+
+    # Load
+    if args.no_jax:
+        flax_params = load_pickle_without_jax(pkl_path)
+    else:
+        flax_params = load_pickle_with_jax(pkl_path)
+
+    # Convert
+    print("Converting weights...")
+    mlx_weights = convert_flax_to_mlx(flax_params, size=args.model)
+
+    # Print summary
+    total_params = sum(w.size for w in mlx_weights.values())
+    print(f"Total parameters: {total_params:,}")
+    print(f"Weight entries: {len(mlx_weights)}")
+
+    # Save
+    output_path = args.output
+    if output_path.endswith('.safetensors'):
+        try:
+            save_safetensors(mlx_weights, output_path)
+        except ImportError:
+            output_path = output_path.replace('.safetensors', '.npz')
+            print("safetensors not installed, saving as npz instead")
+            save_npz(mlx_weights, output_path)
+    else:
+        save_npz(mlx_weights, output_path)
+
+
+if __name__ == '__main__':
+    main()

model.py

@@ -0,0 +1,227 @@
+"""MLX implementation of Thera super-resolution models (air/pro variants)."""
+
+import math
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+
+# --- Utility functions ---
+
+def make_grid(h, w):
+    """Create coordinate grid in [-0.5, 0.5] with pixel centers."""
+    offset_h = 1.0 / (2 * h)
+    offset_w = 1.0 / (2 * w)
+    ys = np.linspace(-0.5 + offset_h, 0.5 - offset_h, h, dtype=np.float32)
+    xs = np.linspace(-0.5 + offset_w, 0.5 - offset_w, w, dtype=np.float32)
+    grid_y, grid_x = np.meshgrid(ys, xs, indexing='ij')
+    return np.stack([grid_y, grid_x], axis=-1)  # (H, W, 2)
+
+
+def interpolate_nearest(coords, grid):
+    """
+    Nearest-neighbor sampling of a grid at given coordinates.
+    Args:
+        coords: mx.array (B, H, W, 2) coordinates in [-0.5, 0.5]
+        grid: mx.array (B, H', W', C) grid to sample from
+    Returns:
+        mx.array (B, H, W, C)
+    """
+    B, Hp, Wp, C = grid.shape
+    _, H, W, _ = coords.shape
+
+    y = coords[..., 0] * Hp + (Hp - 1) / 2.0
+    x = coords[..., 1] * Wp + (Wp - 1) / 2.0
+
+    y_idx = mx.clip(mx.round(y).astype(mx.int32), 0, Hp - 1)
+    x_idx = mx.clip(mx.round(x).astype(mx.int32), 0, Wp - 1)
+
+    flat_idx = y_idx * Wp + x_idx  # (B, H, W)
+    batch_offset = mx.arange(B).reshape(B, 1, 1) * (Hp * Wp)
+    global_idx = (flat_idx + batch_offset).reshape(-1)  # (B*H*W,)
+
+    grid_flat = grid.reshape(-1, C)  # (B*Hp*Wp, C)
+    result = grid_flat[global_idx]  # (B*H*W, C)
+    return result.reshape(B, H, W, C)
+
+
+# --- RDN Backbone ---
+
+class RDBConv(nn.Module):
+    """Single convolution layer within a Residual Dense Block."""
+    def __init__(self, in_channels: int, growth_rate: int, kernel_size: int = 3):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, growth_rate, kernel_size,
+                              padding=(kernel_size - 1) // 2)
+
+    def __call__(self, x):
+        out = nn.relu(self.conv(x))
+        return mx.concatenate([x, out], axis=-1)
+
+
+class RDB(nn.Module):
+    """Residual Dense Block."""
+    def __init__(self, g0: int, growth_rate: int, n_conv_layers: int):
+        super().__init__()
+        self.convs = [
+            RDBConv(g0 + i * growth_rate, growth_rate)
+            for i in range(n_conv_layers)
+        ]
+        total_ch = g0 + n_conv_layers * growth_rate
+        self.local_fusion = nn.Conv2d(total_ch, g0, kernel_size=1)
+
+    def __call__(self, x):
+        res = x
+        for conv in self.convs:
+            x = conv(x)
+        x = self.local_fusion(x)
+        return x + res
+
+
+class RDN(nn.Module):
+    """Residual Dense Network backbone (config B)."""
+    def __init__(self, n_colors: int = 3, g0: int = 64):
+        super().__init__()
+        D, C, G = 16, 8, 64  # config B
+
+        self.sfe1 = nn.Conv2d(n_colors, g0, kernel_size=3, padding=1)
+        self.sfe2 = nn.Conv2d(g0, g0, kernel_size=3, padding=1)
+        self.rdbs = [RDB(g0, G, C) for _ in range(D)]
+        self.gff_1x1 = nn.Conv2d(D * g0, g0, kernel_size=1)
+        self.gff_3x3 = nn.Conv2d(g0, g0, kernel_size=3, padding=1)
+
+    def __call__(self, x):
+        f1 = self.sfe1(x)
+        x = self.sfe2(f1)
+
+        rdb_outs = []
+        for rdb in self.rdbs:
+            x = rdb(x)
+            rdb_outs.append(x)
+
+        x = mx.concatenate(rdb_outs, axis=-1)
+        x = self.gff_1x1(x)
+        x = self.gff_3x3(x)
+        return x + f1
+
+
+# --- Thera Model ---
+
+class Thera(nn.Module):
+    """
+    Thera: arbitrary-scale super-resolution using neural heat fields.
+
+    Stages:
+    1. Encoder (RDN backbone) produces features at source resolution
+    2. Optional refinement tail (identity for air, SwinIR for pro)
+    3. Hypernetwork (1x1 conv) predicts per-pixel field parameters
+    4. Heat field decoder produces RGB residuals
+    """
+    OUT_DIM = 3
+    W0 = 1.0
+    MEAN = np.array([0.4488, 0.4371, 0.4040], dtype=np.float32)
+    VAR = np.array([0.25, 0.25, 0.25], dtype=np.float32)
+
+    def __init__(self, size='air'):
+        super().__init__()
+        self.size = size
+        self.hidden_dim = 32 if size == 'air' else 512
+
+        # Field params: Dense kernel + Thermal phase (alphabetical order)
+        n_field_params = self.hidden_dim * self.OUT_DIM + self.hidden_dim
+
+        self.encoder = RDN(n_colors=3, g0=64)
+
+        # Refinement tail
+        if size == 'pro':
+            from swin_ir import SwinIRTail
+            self.refine = SwinIRTail(
+                in_channels=64, embed_dim=180,
+                depths=(7, 6), num_heads=(6, 6),
+                window_size=8, mlp_ratio=2.0, num_feat=64)
+        # For 'air', no refine module (identity)
+
+        self.out_conv = nn.Conv2d(64, n_field_params, kernel_size=1)
+
+        self.k = mx.array(0.0)
+        self.components = mx.zeros((2, self.hidden_dim))
+
+    def encode(self, source_norm):
+        """Run encoder + optional refinement tail."""
+        x = self.encoder(source_norm)
+        if self.size == 'pro':
+            x = self.refine(x)
+        return x
+
+    def decode(self, encoding, target_coords, t):
+        """Predict RGB residuals at target coordinates."""
+        sampled = interpolate_nearest(target_coords, encoding)
+        phi = self.out_conv(sampled)
+
+        hd = self.hidden_dim
+        kernel = phi[..., :hd * self.OUT_DIM].reshape(
+            *phi.shape[:-1], hd, self.OUT_DIM)
+        phase = phi[..., hd * self.OUT_DIM:]
+
+        Hs, Ws = encoding.shape[1], encoding.shape[2]
+        source_grid = mx.array(make_grid(Hs, Ws))
+        source_coords = mx.broadcast_to(
+            source_grid[None], (encoding.shape[0],) + source_grid.shape)
+        nearest_src = interpolate_nearest(target_coords, source_coords)
+        rel_coords = target_coords - nearest_src
+        rel_coords_scaled = mx.concatenate([
+            rel_coords[..., 0:1] * Hs,
+            rel_coords[..., 1:2] * Ws,
+        ], axis=-1)
+
+        x = rel_coords_scaled @ self.components
+        norm = mx.linalg.norm(self.components, axis=0)
+        t_4d = t[:, :, None, None] if t.ndim == 2 else t.reshape(-1, 1, 1, 1)
+        decay = mx.exp(-((self.W0 * norm) ** 2) * self.k * t_4d)
+        x = mx.sin(self.W0 * x + phase) * decay
+        out = mx.sum(x[..., None] * kernel, axis=-2)
+
+        return out
+
+    def upscale(self, source, target_h, target_w, ensemble=False):
+        mean = mx.array(self.MEAN)
+        var = mx.array(self.VAR)
+        std = mx.sqrt(var)
+
+        if ensemble:
+            outs = []
+            for k_rot in range(4):
+                src = mx.array(np.rot90(np.array(source), k=k_rot))
+                th = target_w if k_rot % 2 else target_h
+                tw = target_h if k_rot % 2 else target_w
+                out = self._upscale_single(src, th, tw, mean, var, std)
+                mx.eval(out)
+                out_np = np.rot90(np.array(out), k=-k_rot)
+                outs.append(out_np)
+            result = np.stack(outs).mean(0).clip(0.0, 1.0)
+            return mx.array((result * 255).round().astype(np.uint8))
+        else:
+            out = self._upscale_single(source, target_h, target_w, mean, var, std)
+            out = mx.clip(out, 0.0, 1.0)
+            return (out * 255 + 0.5).astype(mx.uint8)
+
+    def _upscale_single(self, source, target_h, target_w, mean, var, std):
+        Hs, Ws = source.shape[0], source.shape[1]
+        t = mx.array([(target_h / Hs) ** -2], dtype=mx.float32)[None]
+
+        target_grid = mx.array(make_grid(target_h, target_w))[None]
+        source_4d = source[None]
+        source_up = interpolate_nearest(target_grid, source_4d)
+
+        source_norm = (source_4d - mean) / std
+        encoding = self.encode(source_norm)
+
+        coords = mx.array(make_grid(target_h, target_w))[None]
+        residual = self.decode(encoding, coords, t)
+
+        out = residual * std + mean + source_up
+        return out[0]
+
+
+# Backwards compatibility alias
+TheraRDNAir = Thera

requirements.txt

@@ -0,0 +1,25 @@
+# Thera MLX — requirements
+# Install with: pip install -r requirements.txt
+#
+# Requires: Python 3.10+, macOS with Apple Silicon (M1/M2/M3/M4)
+
+# Core ML
+mlx>=0.16.0
+
+# Image I/O
+Pillow>=10.0.0
+numpy>=1.24.0
+
+# Web UI
+flask>=3.0.0
+
+# Weight loading
+safetensors>=0.4.0
+
+# Video support (optional — needed for the Video tab)
+# Installs a bundled ffmpeg so you don't need a system install
+imageio[ffmpeg]>=2.34.0
+
+# Weight conversion from Flax (only needed if re-converting from source)
+# jax[cpu]>=0.4.0        # uncomment if you have JAX available
+# flax>=0.8.0            # uncomment if you have JAX available

run.py

@@ -0,0 +1,132 @@
+#!/usr/bin/env python3
+"""
+Thera MLX — Arbitrary-scale super-resolution on Apple Silicon.
+
+Usage:
+    python run.py                         # Launch web UI (default port 5005)
+    python run.py --port 8080             # Custom port
+    python run.py run input.png out.png --scale 4 --model pro
+    python run.py run input.png out.png --scale 2 --model air --ensemble --tiles 2
+    python run.py convert --model air     # Download + convert weights
+"""
+
+import argparse
+import sys
+import os
+
+# Make sure this directory is on the path so all modules import cleanly
+sys.path.insert(0, os.path.dirname(__file__))
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        prog="thera-mlx",
+        description="Thera MLX — arbitrary-scale super-resolution on Apple Silicon",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__,
+    )
+    sub = parser.add_subparsers(dest="command")
+
+    # ── gui (default) ────────────────────────────────────────────────────────
+    gui_p = sub.add_parser("gui", help="Launch web UI (default)")
+    gui_p.add_argument("--port", type=int, default=5005, help="Port (default: 5005)")
+    gui_p.add_argument("--host", type=str, default="127.0.0.1")
+
+    # ── run ──────────────────────────────────────────────────────────────────
+    run_p = sub.add_parser("run", help="Upscale a single image from the CLI")
+    run_p.add_argument("input",  help="Input image path")
+    run_p.add_argument("output", help="Output image path")
+    scale_g = run_p.add_mutually_exclusive_group(required=True)
+    scale_g.add_argument("--scale", type=float, help="Scale factor  (e.g. 2, 4)")
+    scale_g.add_argument("--size",  type=int, nargs=2, metavar=("H", "W"),
+                          help="Exact target size  height width")
+    run_p.add_argument("--model",    choices=["air", "pro"], default="air",
+                        help="Model variant: air (fast) or pro (quality). Default: air")
+    run_p.add_argument("--weights",  default=None,
+                        help="Path to weights file (optional — uses bundled weights by default)")
+    run_p.add_argument("--ensemble", action="store_true",
+                        help="Geometric self-ensemble (4 rotations) — higher quality, slower")
+    run_p.add_argument("--tiles",    type=int, choices=[2, 3, 4], default=None,
+                        help="Tile NxN to reduce RAM usage (2, 3, or 4)")
+
+    # ── convert ──────────────────────────────────────────────────────────────
+    conv_p = sub.add_parser("convert", help="Download + convert weights from Flax format")
+    conv_p.add_argument("--model",  choices=["air", "pro"], default="air")
+    conv_p.add_argument("--input",  default=None,
+                         help="Path to model.pkl (downloads from HuggingFace if omitted)")
+    conv_p.add_argument("--no-jax", action="store_true",
+                         help="Load pickle without JAX installed")
+
+    args = parser.parse_args()
+
+    # Default: no subcommand → launch GUI
+    if args.command is None or args.command == "gui":
+        _cmd_gui(args if args.command else argparse.Namespace(host="127.0.0.1", port=5005))
+    elif args.command == "run":
+        _cmd_run(args)
+    elif args.command == "convert":
+        _cmd_convert(args)
+
+
+# ── Command implementations ───────────────────────────────────────────────────
+
+def _cmd_gui(args):
+    import logging
+    logging.getLogger("werkzeug").setLevel(logging.ERROR)
+    from ui import app
+    print(f"\n  Thera MLX  →  http://{args.host}:{args.port}\n")
+    app.run(host=args.host, port=args.port, debug=False, threaded=True)
+
+
+def _cmd_run(args):
+    from upscale import upscale_file
+    upscale_file(
+        args.input, args.output,
+        scale=args.scale, size=args.size,
+        model_size=args.model, weights_path=args.weights,
+        ensemble=args.ensemble, tiles=args.tiles,
+    )
+
+
+def _cmd_convert(args):
+    from pathlib import Path
+    from upscale import WEIGHTS_DIR
+    from convert import (
+        REPO_IDS, download_model, convert_flax_to_mlx,
+        load_pickle_with_jax, load_pickle_without_jax,
+        save_safetensors, save_npz,
+    )
+
+    output_path = str(WEIGHTS_DIR / f"weights-{args.model}.safetensors")
+
+    if args.input is None:
+        repo = REPO_IDS[args.model]
+        print(f"Downloading model from HuggingFace ({repo})...")
+        pkl_path = download_model(args.model)
+    else:
+        pkl_path = args.input
+
+    flax_params = (load_pickle_without_jax if args.no_jax else load_pickle_with_jax)(pkl_path)
+
+    print("Converting weights...")
+    mlx_weights = convert_flax_to_mlx(flax_params, size=args.model)
+
+    total_params = sum(w.size for w in mlx_weights.values())
+    print(f"  Parameters : {total_params:,}")
+    print(f"  Weight keys: {len(mlx_weights)}")
+
+    if output_path.endswith(".safetensors"):
+        try:
+            save_safetensors(mlx_weights, output_path)
+        except ImportError:
+            output_path = output_path.replace(".safetensors", ".npz")
+            print("safetensors not installed — saving as .npz instead")
+            save_npz(mlx_weights, output_path)
+    else:
+        save_npz(mlx_weights, output_path)
+
+    print(f"Saved → {output_path}")
+
+
+if __name__ == "__main__":
+    main()

swin_ir.py

@@ -0,0 +1,286 @@
+"""MLX implementation of SwinIR (Swin Transformer for Image Restoration)."""
+
+import math
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+
+def window_partition(x, window_size):
+    """Partition into non-overlapping windows.
+    x: (B, H, W, C) → (B * nH * nW, ws, ws, C)
+    """
+    B, H, W, C = x.shape
+    nH, nW = H // window_size, W // window_size
+    x = x.reshape(B, nH, window_size, nW, window_size, C)
+    x = x.transpose(0, 1, 3, 2, 4, 5)  # B, nH, nW, ws, ws, C
+    return x.reshape(-1, window_size, window_size, C)
+
+
+def window_reverse(windows, window_size, H, W):
+    """Reverse window_partition.
+    windows: (B * nH * nW, ws, ws, C) → (B, H, W, C)
+    """
+    nH, nW = H // window_size, W // window_size
+    B = windows.shape[0] // (nH * nW)
+    x = windows.reshape(B, nH, nW, window_size, window_size, -1)
+    x = x.transpose(0, 1, 3, 2, 4, 5)  # B, nH, ws, nW, ws, C
+    return x.reshape(B, H, W, -1)
+
+
+def make_attn_mask(shift_size, window_size, H, W):
+    """Create attention mask for shifted window self-attention."""
+    if shift_size == 0:
+        return None
+
+    mask = np.zeros((1, H, W, 1))
+    h_slices = (slice(0, -window_size),
+                slice(-window_size, -shift_size),
+                slice(-shift_size, None))
+    w_slices = (slice(0, -window_size),
+                slice(-window_size, -shift_size),
+                slice(-shift_size, None))
+
+    cnt = 0
+    for h in h_slices:
+        for w in w_slices:
+            mask[:, h, w, :] = cnt
+            cnt += 1
+
+    mask_windows = window_partition(mx.array(mask), window_size)  # nW, ws, ws, 1
+    mask_windows = mask_windows.reshape(-1, window_size * window_size)  # nW, ws*ws
+    attn_mask = mask_windows[:, :, None] - mask_windows[:, None, :]  # nW, ws*ws, ws*ws
+    attn_mask = mx.where(attn_mask != 0, -100.0, 0.0)
+    return attn_mask
+
+
+def make_relative_position_index(window_size):
+    """Compute relative position index for each pair in a window."""
+    coords_h = np.arange(window_size)
+    coords_w = np.arange(window_size)
+    # Note: meshgrid with indexing='ij' matches the Flax implementation
+    coords = np.stack(np.meshgrid(coords_w, coords_h, indexing='ij'))  # 2, ws, ws
+    coords_flat = coords.reshape(2, -1)  # 2, ws*ws
+    relative = coords_flat[:, :, None] - coords_flat[:, None, :]  # 2, ws*ws, ws*ws
+    relative = relative.transpose(1, 2, 0)  # ws*ws, ws*ws, 2
+    relative[:, :, 0] += window_size - 1
+    relative[:, :, 1] += window_size - 1
+    relative[:, :, 0] *= 2 * window_size - 1
+    return relative.sum(-1)  # ws*ws, ws*ws
+
+
+class MLP(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.fc1 = nn.Linear(dim, hidden_dim)
+        self.fc2 = nn.Linear(hidden_dim, dim)
+
+    def __call__(self, x):
+        return self.fc2(nn.gelu(self.fc1(x)))
+
+
+class WindowAttention(nn.Module):
+    """Window-based multi-head self-attention with relative position bias."""
+
+    def __init__(self, dim, window_size, num_heads):
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3)
+        self.proj = nn.Linear(dim, dim)
+
+        # Relative position bias table: (2*ws-1)^2 entries, num_heads
+        num_rel = (2 * window_size - 1) ** 2
+        self.relative_position_bias_table = mx.zeros((num_rel, num_heads))
+
+        # Precompute relative position index (as flat int array for indexing)
+        idx = make_relative_position_index(window_size)
+        self._rel_pos_index = mx.array(idx.reshape(-1), dtype=mx.int32)
+
+    def __call__(self, x, mask=None):
+        B, N, C = x.shape  # N = ws*ws
+
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, self.head_dim)
+        qkv = qkv.transpose(2, 0, 3, 1, 4)  # 3, B, heads, N, head_dim
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        q = q * self.scale
+        attn = q @ k.transpose(0, 1, 3, 2)  # B, heads, N, N
+
+        # Add relative position bias
+        rel_pos_bias = self.relative_position_bias_table[
+            self._rel_pos_index
+        ].reshape(N, N, -1)
+        rel_pos_bias = rel_pos_bias.transpose(2, 0, 1)  # heads, N, N
+        attn = attn + rel_pos_bias[None]
+
+        # Apply mask for shifted windows
+        if mask is not None:
+            nW = mask.shape[0]
+            attn = attn.reshape(B // nW, nW, self.num_heads, N, N)
+            attn = attn + mask[None, :, None, :, :]
+            attn = attn.reshape(-1, self.num_heads, N, N)
+
+        attn = mx.softmax(attn, axis=-1)
+        x = (attn @ v).transpose(0, 2, 1, 3).reshape(B, N, C)
+        return self.proj(x)
+
+
+class SwinTransformerBlock(nn.Module):
+    """Swin Transformer Block with optional shifted windows."""
+
+    def __init__(self, dim, num_heads, window_size, shift_size, mlp_ratio=2.0):
+        super().__init__()
+        self.dim = dim
+        self.window_size = window_size
+        self.shift_size = shift_size
+
+        self.norm1 = nn.LayerNorm(dim, eps=1e-5)
+        self.attn = WindowAttention(dim, window_size, num_heads)
+        self.norm2 = nn.LayerNorm(dim, eps=1e-5)
+        self.mlp = MLP(dim, int(dim * mlp_ratio))
+
+    def __call__(self, x, H, W):
+        B, L, C = x.shape
+        shortcut = x
+
+        x = self.norm1(x)
+        x = x.reshape(B, H, W, C)
+
+        ws = self.window_size
+        shift = self.shift_size
+
+        # Cyclic shift
+        if shift > 0:
+            x = mx.roll(x, (-shift, -shift), axis=(1, 2))
+
+        # Partition windows
+        x_windows = window_partition(x, ws)  # nW*B, ws, ws, C
+        x_windows = x_windows.reshape(-1, ws * ws, C)
+
+        # Window attention with mask
+        mask = make_attn_mask(shift, ws, H, W) if shift > 0 else None
+        attn_windows = self.attn(x_windows, mask)
+
+        # Merge windows
+        attn_windows = attn_windows.reshape(-1, ws, ws, C)
+        x = window_reverse(attn_windows, ws, H, W)
+
+        # Reverse cyclic shift
+        if shift > 0:
+            x = mx.roll(x, (shift, shift), axis=(1, 2))
+
+        x = x.reshape(B, H * W, C)
+
+        # Residual + FFN
+        x = shortcut + x
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class RSTB(nn.Module):
+    """Residual Swin Transformer Block."""
+
+    def __init__(self, dim, depth, num_heads, window_size, mlp_ratio=2.0):
+        super().__init__()
+        self.blocks = []
+        for i in range(depth):
+            shift = 0 if (i % 2 == 0) else window_size // 2
+            self.blocks.append(
+                SwinTransformerBlock(dim, num_heads, window_size, shift, mlp_ratio)
+            )
+        self.conv = nn.Conv2d(dim, dim, kernel_size=3, padding=1)
+
+    def __call__(self, x, H, W):
+        res = x
+        for block in self.blocks:
+            x = block(x, H, W)
+
+        # PatchUnEmbed → Conv → PatchEmbed
+        B, L, C = x.shape
+        x = x.reshape(B, H, W, C)
+        x = self.conv(x)
+        x = x.reshape(B, H * W, C)
+        return x + res
+
+
+class SwinIRTail(nn.Module):
+    """SwinIR tail module for Thera-pro.
+
+    Takes 64-channel RDN output and refines it, outputting 64 channels.
+    """
+
+    def __init__(self, in_channels=64, embed_dim=180, depths=(7, 6),
+                 num_heads=(6, 6), window_size=8, mlp_ratio=2.0, num_feat=64):
+        super().__init__()
+        self.window_size = window_size
+        self.embed_dim = embed_dim
+
+        # conv_first
+        self.conv_first = nn.Conv2d(in_channels, embed_dim, kernel_size=3, padding=1)
+
+        # PatchEmbed norm
+        self.patch_embed_norm = nn.LayerNorm(embed_dim, eps=1e-5)
+
+        # RSTB layers
+        self.layers = [
+            RSTB(embed_dim, depth, nh, window_size, mlp_ratio)
+            for depth, nh in zip(depths, num_heads)
+        ]
+
+        # Final norm
+        self.norm = nn.LayerNorm(embed_dim, eps=1e-5)
+
+        # conv_after_body
+        self.conv_after_body = nn.Conv2d(embed_dim, embed_dim, kernel_size=3, padding=1)
+
+        # conv_before_upsample → output 64 channels
+        self.conv_last = nn.Conv2d(embed_dim, num_feat, kernel_size=3, padding=1)
+
+    def __call__(self, x):
+        _, h_orig, w_orig, _ = x.shape
+
+        # Pad to multiple of window_size
+        ws = self.window_size
+        pad_h = (ws - h_orig % ws) % ws
+        pad_w = (ws - w_orig % ws) % ws
+        if pad_h > 0 or pad_w > 0:
+            # Reflect padding
+            x = mx.pad(x, [(0, 0), (0, pad_h), (0, pad_w), (0, 0)], mode='edge')
+
+        _, H, W, _ = x.shape
+
+        # conv_first
+        x = self.conv_first(x)
+        res = x
+
+        # PatchEmbed (reshape + norm)
+        B = x.shape[0]
+        x = x.reshape(B, H * W, self.embed_dim)
+        x = self.patch_embed_norm(x)
+
+        # RSTB layers
+        for layer in self.layers:
+            x = layer(x, H, W)
+
+        # Final norm
+        x = self.norm(x)
+
+        # PatchUnEmbed
+        x = x.reshape(B, H, W, self.embed_dim)
+
+        # conv_after_body + residual
+        x = self.conv_after_body(x) + res
+
+        # conv_before_upsample
+        x = nn.leaky_relu(self.conv_last(x))
+
+        # Remove padding
+        if pad_h > 0 or pad_w > 0:
+            x = x[:, :h_orig, :w_orig, :]
+
+        return x

ui.py

@@ -0,0 +1,1435 @@
+#!/usr/bin/env python3
+"""
+Web UI for Thera MLX super-resolution.
+
+Supports single image, batch, and video upscaling.
+
+Usage:
+    python3 ui.py
+    python3 ui.py --port 8080
+"""
+
+import argparse
+import glob
+import io
+import json
+import os
+import shutil
+import subprocess
+import tempfile
+import threading
+import time
+import uuid
+import zipfile
+
+import mlx.core as mx
+import numpy as np
+from flask import Flask, request, jsonify, send_file, Response
+from PIL import Image
+
+from model import Thera
+
+
+def _find_ffmpeg():
+    """Find ffmpeg binary — system PATH first, then imageio_ffmpeg fallback."""
+    path = shutil.which("ffmpeg")
+    if path:
+        return path
+    try:
+        import imageio_ffmpeg
+        return imageio_ffmpeg.get_ffmpeg_exe()
+    except ImportError:
+        return None
+
+
+def _find_ffprobe():
+    """Find ffprobe binary — system PATH first, then imageio_ffmpeg fallback."""
+    path = shutil.which("ffprobe")
+    if path:
+        return path
+    try:
+        import imageio_ffmpeg
+        ff = imageio_ffmpeg.get_ffmpeg_exe()
+        # ffprobe is in the same directory
+        probe = os.path.join(os.path.dirname(ff),
+                             ff.replace("ffmpeg", "ffprobe").split("/")[-1])
+        if os.path.exists(probe):
+            return probe
+        # some builds bundle it as ffprobe next to ffmpeg
+        probe2 = ff.replace("ffmpeg", "ffprobe")
+        if os.path.exists(probe2):
+            return probe2
+    except ImportError:
+        pass
+    return None
+
+
+FFMPEG = _find_ffmpeg()
+FFPROBE = _find_ffprobe()
+
+# ---------------------------------------------------------------------------
+# App setup
+# ---------------------------------------------------------------------------
+
+app = Flask(__name__)
+UPLOAD_DIR = tempfile.mkdtemp(prefix="thera_ui_")
+
+# Job tracking for async video processing
+_jobs = {}
+_jobs_lock = threading.Lock()
+
+# ---------------------------------------------------------------------------
+# Model cache
+# ---------------------------------------------------------------------------
+
+_model_cache = {}
+
+
+def get_model(size):
+    if size not in _model_cache:
+        from upscale import load_weights
+        model = Thera(size=size)
+        weights_dir = os.path.join(os.path.dirname(__file__), "weights")
+        weights_path = os.path.join(weights_dir, f"weights-{size}.safetensors")
+        if not os.path.exists(weights_path):
+            raise FileNotFoundError(
+                f"Weights not found: {weights_path}\n"
+                f"Run: python3 convert.py --model {size}")
+        model = load_weights(model, weights_path)
+        mx.eval(model.parameters())
+        _model_cache[size] = model
+    return _model_cache[size]
+
+
+# ---------------------------------------------------------------------------
+# Core upscale
+# ---------------------------------------------------------------------------
+
+def upscale_image(img_np, scale, model_size, ensemble, tiles=None):
+    model = get_model(model_size)
+    h, w = img_np.shape[:2]
+    th, tw = round(h * scale), round(w * scale)
+    source_f = img_np.astype(np.float32) / 255.0
+
+    if tiles and tiles > 1:
+        from upscale import upscale_tiled
+        return upscale_tiled(model, source_f, th, tw, tiles, ensemble=ensemble)
+    else:
+        source = mx.array(source_f)
+        result = model.upscale(source, th, tw, ensemble=ensemble)
+        mx.eval(result)
+        return np.array(result)
+
+
+# ---------------------------------------------------------------------------
+# API routes
+# ---------------------------------------------------------------------------
+
+@app.route("/api/upscale", methods=["POST"])
+def api_upscale():
+    f = request.files.get("image")
+    if not f:
+        return jsonify(error="No image uploaded"), 400
+
+    scale = float(request.form.get("scale", 2.0))
+    model_size = request.form.get("model", "air")
+    ensemble = request.form.get("ensemble", "false") == "true"
+    tiles_str = request.form.get("tiles", "1")
+    tiles = int(tiles_str) if tiles_str.isdigit() and int(tiles_str) > 1 else None
+
+    img = np.array(Image.open(f.stream).convert("RGB"))
+    t0 = time.perf_counter()
+    result = upscale_image(img, scale, model_size, ensemble, tiles=tiles)
+    elapsed = time.perf_counter() - t0
+
+    buf = io.BytesIO()
+    Image.fromarray(result).save(buf, format="PNG")
+    buf.seek(0)
+
+    h, w = img.shape[:2]
+    th, tw = result.shape[:2]
+
+    resp = send_file(buf, mimetype="image/png", download_name="upscaled.png")
+    resp.headers["X-Info"] = json.dumps({
+        "src": f"{w}x{h}", "dst": f"{tw}x{th}",
+        "scale": scale, "model": model_size, "time": round(elapsed, 1)
+    })
+    return resp
+
+
+@app.route("/api/batch", methods=["POST"])
+def api_batch():
+    files = request.files.getlist("images")
+    if not files:
+        return jsonify(error="No images uploaded"), 400
+
+    scale = float(request.form.get("scale", 2.0))
+    model_size = request.form.get("model", "air")
+    ensemble = request.form.get("ensemble", "false") == "true"
+
+    buf = io.BytesIO()
+    t0 = time.perf_counter()
+    with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
+        for f in files:
+            img = np.array(Image.open(f.stream).convert("RGB"))
+            result = upscale_image(img, scale, model_size, ensemble)
+            img_buf = io.BytesIO()
+            Image.fromarray(result).save(img_buf, format="PNG")
+            name = os.path.splitext(f.filename)[0] + f"_thera_{scale}x.png"
+            zf.writestr(name, img_buf.getvalue())
+
+    elapsed = time.perf_counter() - t0
+    buf.seek(0)
+
+    resp = send_file(buf, mimetype="application/zip",
+                     download_name="thera_batch.zip")
+    resp.headers["X-Info"] = json.dumps({
+        "count": len(files), "scale": scale,
+        "model": model_size, "time": round(elapsed, 1)
+    })
+    return resp
+
+
+# ---------------------------------------------------------------------------
+# Video processing (async with progress)
+# ---------------------------------------------------------------------------
+
+def get_video_info(path):
+    """Get video metadata using ffprobe if available, else parse ffmpeg stderr."""
+    if FFPROBE:
+        cmd = [
+            FFPROBE, "-v", "quiet", "-print_format", "json",
+            "-show_streams", "-show_format", path
+        ]
+        result = subprocess.run(cmd, capture_output=True, text=True)
+        info = json.loads(result.stdout)
+        stream = next(s for s in info["streams"] if s["codec_type"] == "video")
+        fps_parts = stream["r_frame_rate"].split("/")
+        fps = float(fps_parts[0]) / float(fps_parts[1])
+        w, h = int(stream["width"]), int(stream["height"])
+        duration = float(info["format"].get("duration", 0))
+        return fps, w, h, duration
+    else:
+        # Fallback: parse ffmpeg -i stderr
+        import re
+        result = subprocess.run(
+            [FFMPEG, "-i", path], capture_output=True, text=True)
+        stderr = result.stderr
+        # Parse "Duration: 00:00:10.00"
+        dur_m = re.search(r"Duration:\s*(\d+):(\d+):(\d+\.\d+)", stderr)
+        duration = 0.0
+        if dur_m:
+            duration = int(dur_m[1]) * 3600 + int(dur_m[2]) * 60 + float(dur_m[3])
+        # Parse "1920x1080" and fps
+        vid_m = re.search(r"(\d{2,5})x(\d{2,5})", stderr)
+        w, h = (int(vid_m[1]), int(vid_m[2])) if vid_m else (0, 0)
+        fps_m = re.search(r"(\d+(?:\.\d+)?)\s*fps", stderr)
+        fps = float(fps_m[1]) if fps_m else 30.0
+        return fps, w, h, duration
+
+
+def has_audio(path):
+    """Check if video has an audio stream."""
+    if FFPROBE:
+        probe = subprocess.run(
+            [FFPROBE, "-v", "quiet", "-select_streams", "a",
+             "-show_entries", "stream=codec_type", path],
+            capture_output=True, text=True)
+        return "audio" in probe.stdout
+    else:
+        result = subprocess.run(
+            [FFMPEG, "-i", path], capture_output=True, text=True)
+        return "Audio:" in result.stderr
+
+
+def video_worker(job_id, video_path, scale, model_size):
+    job = _jobs[job_id]
+    try:
+        job["status"] = "analyzing"
+        fps, src_w, src_h, duration = get_video_info(video_path)
+        tw, th = round(src_w * scale), round(src_h * scale)
+        job["src"] = f"{src_w}x{src_h}"
+        job["dst"] = f"{tw}x{th}"
+
+        tmpdir = tempfile.mkdtemp(prefix="thera_vid_")
+        frames_dir = os.path.join(tmpdir, "frames")
+        upscaled_dir = os.path.join(tmpdir, "upscaled")
+        os.makedirs(frames_dir)
+        os.makedirs(upscaled_dir)
+
+        # Extract frames
+        job["status"] = "extracting"
+        subprocess.run([
+            FFMPEG, "-i", video_path, "-vsync", "0",
+            os.path.join(frames_dir, "frame_%06d.png")
+        ], capture_output=True, check=True)
+
+        frame_files = sorted(glob.glob(os.path.join(frames_dir, "frame_*.png")))
+        total = len(frame_files)
+        job["total_frames"] = total
+
+        # Upscale frames
+        job["status"] = "upscaling"
+        model = get_model(model_size)
+        t0 = time.perf_counter()
+
+        for i, frame_path in enumerate(frame_files):
+            img = np.array(Image.open(frame_path).convert("RGB"))
+            source = mx.array(img.astype(np.float32) / 255.0)
+            fh, fw = img.shape[:2]
+            result = model.upscale(source, round(fh * scale), round(fw * scale))
+            mx.eval(result)
+            out_path = os.path.join(upscaled_dir, os.path.basename(frame_path))
+            Image.fromarray(np.array(result)).save(out_path)
+
+            elapsed = time.perf_counter() - t0
+            eta = (elapsed / (i + 1)) * (total - i - 1) if i > 0 else 0
+            job["current_frame"] = i + 1
+            job["eta"] = round(eta)
+            job["fps"] = round((i + 1) / elapsed, 1) if elapsed > 0 else 0
+
+        # Encode
+        job["status"] = "encoding"
+        output_path = os.path.join(tmpdir, "upscaled.mp4")
+
+        ffmpeg_cmd = [
+            FFMPEG, "-y",
+            "-framerate", str(fps),
+            "-i", os.path.join(upscaled_dir, "frame_%06d.png"),
+        ]
+
+        # Check for audio
+        audio = has_audio(video_path)
+
+        if audio:
+            ffmpeg_cmd += ["-i", video_path, "-map", "0:v", "-map", "1:a",
+                           "-shortest"]
+
+        ffmpeg_cmd += [
+            "-c:v", "libx264", "-preset", "medium", "-crf", "18",
+            "-pix_fmt", "yuv420p",
+        ]
+        if audio:
+            ffmpeg_cmd += ["-c:a", "aac", "-b:a", "192k"]
+
+        ffmpeg_cmd.append(output_path)
+        subprocess.run(ffmpeg_cmd, capture_output=True, check=True)
+
+        total_time = time.perf_counter() - t0
+        job["status"] = "done"
+        job["output_path"] = output_path
+        job["time"] = round(total_time, 1)
+        job["tmpdir"] = tmpdir
+
+    except Exception as e:
+        job["status"] = "error"
+        job["error"] = str(e)
+
+
+@app.route("/api/video/start", methods=["POST"])
+def api_video_start():
+    f = request.files.get("video")
+    if not f:
+        return jsonify(error="No video uploaded"), 400
+
+    if not FFMPEG:
+        return jsonify(error="ffmpeg not found. Install with: pip3 install imageio[ffmpeg]"), 400
+
+    scale = float(request.form.get("scale", 2.0))
+    model_size = request.form.get("model", "air")
+
+    job_id = str(uuid.uuid4())[:8]
+    video_path = os.path.join(UPLOAD_DIR, f"{job_id}_input.mp4")
+    f.save(video_path)
+
+    with _jobs_lock:
+        _jobs[job_id] = {
+            "status": "queued", "current_frame": 0, "total_frames": 0,
+            "eta": 0, "fps": 0, "scale": scale, "model": model_size,
+        }
+
+    t = threading.Thread(target=video_worker,
+                         args=(job_id, video_path, scale, model_size),
+                         daemon=True)
+    t.start()
+
+    return jsonify(job_id=job_id)
+
+
+@app.route("/api/video/progress/<job_id>")
+def api_video_progress(job_id):
+    job = _jobs.get(job_id)
+    if not job:
+        return jsonify(error="Unknown job"), 404
+    safe = {k: v for k, v in job.items()
+            if k not in ("output_path", "tmpdir")}
+    return jsonify(safe)
+
+
+@app.route("/api/video/download/<job_id>")
+def api_video_download(job_id):
+    job = _jobs.get(job_id)
+    if not job or job["status"] != "done":
+        return jsonify(error="Not ready"), 400
+    return send_file(job["output_path"], mimetype="video/mp4",
+                     download_name="thera_upscaled.mp4")
+
+
+# ---------------------------------------------------------------------------
+# Frontend
+# ---------------------------------------------------------------------------
+
+@app.route("/")
+def index():
+    return HTML_PAGE
+
+
+HTML_PAGE = r"""<!DOCTYPE html>
+<html lang="en">
+<head>
+<meta charset="UTF-8">
+<meta name="viewport" content="width=device-width, initial-scale=1.0">
+<title>Thera MLX</title>
+<style>
+:root {
+  --bg: #0f0f0f;
+  --surface: #1a1a1a;
+  --surface2: #242424;
+  --border: #333;
+  --text: #e8e8e8;
+  --text2: #999;
+  --accent: #6c63ff;
+  --accent-hover: #7c74ff;
+  --green: #4caf50;
+  --radius: 12px;
+}
+
+* { box-sizing: border-box; margin: 0; padding: 0; }
+
+body {
+  font-family: -apple-system, BlinkMacSystemFont, 'SF Pro Text', system-ui, sans-serif;
+  background: var(--bg);
+  color: var(--text);
+  min-height: 100vh;
+}
+
+.container { max-width: 960px; margin: 0 auto; padding: 24px 20px; }
+
+header {
+  text-align: center;
+  padding: 32px 0 24px;
+}
+header h1 {
+  font-size: 28px;
+  font-weight: 700;
+  letter-spacing: -0.5px;
+}
+header h1 span { color: var(--accent); }
+header p { color: var(--text2); margin-top: 4px; font-size: 14px; }
+
+/* Tabs */
+.tabs {
+  display: flex;
+  gap: 4px;
+  background: var(--surface);
+  border-radius: var(--radius);
+  padding: 4px;
+  margin-bottom: 20px;
+}
+.tab {
+  flex: 1;
+  padding: 10px 16px;
+  border: none;
+  background: none;
+  color: var(--text2);
+  font-size: 14px;
+  font-weight: 500;
+  cursor: pointer;
+  border-radius: 8px;
+  transition: all 0.2s;
+}
+.tab:hover { color: var(--text); }
+.tab.active { background: var(--accent); color: #fff; }
+
+.tab-content { display: none; }
+.tab-content.active { display: block; }
+
+/* Controls */
+.controls {
+  background: var(--surface);
+  border-radius: var(--radius);
+  padding: 20px;
+  margin-bottom: 16px;
+}
+.control-row {
+  display: flex;
+  gap: 16px;
+  align-items: center;
+  flex-wrap: wrap;
+}
+.control-group {
+  display: flex;
+  flex-direction: column;
+  gap: 6px;
+}
+.control-group label {
+  font-size: 12px;
+  font-weight: 600;
+  color: var(--text2);
+  text-transform: uppercase;
+  letter-spacing: 0.5px;
+}
+.control-group select,
+.control-group input[type="range"] {
+  background: var(--surface2);
+  border: 1px solid var(--border);
+  color: var(--text);
+  border-radius: 8px;
+  padding: 8px 12px;
+  font-size: 14px;
+}
+.control-group select { min-width: 100px; cursor: pointer; }
+.control-group input[type="range"] { width: 160px; accent-color: var(--accent); }
+
+.scale-display {
+  font-size: 20px;
+  font-weight: 700;
+  color: var(--accent);
+  min-width: 40px;
+  text-align: center;
+}
+
+.checkbox-label {
+  display: flex;
+  align-items: center;
+  gap: 8px;
+  font-size: 14px;
+  cursor: pointer;
+  color: var(--text2);
+}
+.checkbox-label input { accent-color: var(--accent); }
+
+/* Drop zone */
+.dropzone {
+  border: 2px dashed var(--border);
+  border-radius: var(--radius);
+  padding: 48px 24px;
+  text-align: center;
+  cursor: pointer;
+  transition: all 0.2s;
+  background: var(--surface);
+  margin-bottom: 16px;
+}
+.dropzone:hover, .dropzone.dragover {
+  border-color: var(--accent);
+  background: rgba(108, 99, 255, 0.05);
+}
+.dropzone.has-file {
+  padding: 16px;
+  border-style: solid;
+  border-color: var(--green);
+}
+.dropzone-icon { font-size: 36px; margin-bottom: 8px; }
+.dropzone-text { color: var(--text2); font-size: 14px; }
+.dropzone-text strong { color: var(--text); }
+
+/* Preview panels */
+.preview-area {
+  display: grid;
+  grid-template-columns: 1fr 1fr;
+  gap: 16px;
+  margin-bottom: 16px;
+}
+.preview-panel {
+  background: var(--surface);
+  border-radius: var(--radius);
+  overflow: hidden;
+  min-height: 200px;
+  position: relative;
+}
+.preview-panel .panel-label {
+  position: absolute;
+  top: 8px;
+  left: 12px;
+  font-size: 11px;
+  font-weight: 600;
+  text-transform: uppercase;
+  color: var(--text2);
+  background: rgba(0,0,0,0.6);
+  padding: 3px 8px;
+  border-radius: 4px;
+  z-index: 1;
+}
+.preview-panel img, .preview-panel video {
+  width: 100%;
+  height: 100%;
+  object-fit: contain;
+  display: block;
+}
+
+/* Compare slider */
+.compare-container {
+  display: none;
+  position: relative;
+  background: var(--surface);
+  border-radius: var(--radius);
+  overflow: hidden;
+  margin-bottom: 16px;
+  cursor: col-resize;
+  user-select: none;
+}
+.compare-container.active { display: block; }
+.compare-container img {
+  width: 100%;
+  display: block;
+}
+.compare-overlay {
+  position: absolute;
+  top: 0;
+  left: 0;
+  height: 100%;
+  overflow: hidden;
+}
+.compare-overlay img {
+  position: absolute;
+  top: 0;
+  left: 0;
+  height: 100%;
+  width: auto;
+  max-width: none;
+}
+.compare-line {
+  position: absolute;
+  top: 0;
+  width: 2px;
+  height: 100%;
+  background: #fff;
+  pointer-events: none;
+  z-index: 2;
+  box-shadow: 0 0 6px rgba(0,0,0,0.5);
+}
+.compare-line::after {
+  content: '';
+  position: absolute;
+  top: 50%;
+  left: 50%;
+  transform: translate(-50%, -50%);
+  width: 32px;
+  height: 32px;
+  background: #fff;
+  border-radius: 50%;
+  box-shadow: 0 2px 8px rgba(0,0,0,0.3);
+}
+.compare-line::before {
+  content: '';
+  position: absolute;
+  top: 50%;
+  left: 50%;
+  transform: translate(-50%, -50%);
+  width: 16px;
+  height: 16px;
+  background: var(--accent);
+  border-radius: 50%;
+  z-index: 1;
+}
+.compare-label {
+  position: absolute;
+  top: 8px;
+  font-size: 11px;
+  font-weight: 600;
+  text-transform: uppercase;
+  color: var(--text2);
+  background: rgba(0,0,0,0.6);
+  padding: 3px 8px;
+  border-radius: 4px;
+  z-index: 3;
+  pointer-events: none;
+}
+.compare-label-before { left: 12px; }
+.compare-label-after { right: 12px; }
+
+/* Buttons */
+.btn {
+  padding: 12px 32px;
+  border: none;
+  border-radius: 8px;
+  font-size: 15px;
+  font-weight: 600;
+  cursor: pointer;
+  transition: all 0.2s;
+  display: inline-flex;
+  align-items: center;
+  gap: 8px;
+}
+.btn-primary {
+  background: var(--accent);
+  color: #fff;
+}
+.btn-primary:hover { background: var(--accent-hover); }
+.btn-primary:disabled {
+  opacity: 0.5;
+  cursor: not-allowed;
+}
+.btn-secondary {
+  background: var(--surface2);
+  color: var(--text);
+  border: 1px solid var(--border);
+}
+.btn-secondary:hover { background: var(--border); }
+
+.action-row {
+  display: flex;
+  gap: 12px;
+  align-items: center;
+  flex-wrap: wrap;
+}
+
+/* Progress */
+.progress-container {
+  display: none;
+  margin-bottom: 16px;
+}
+.progress-container.active { display: block; }
+.progress-bar-bg {
+  background: var(--surface2);
+  border-radius: 8px;
+  height: 8px;
+  overflow: hidden;
+  margin-bottom: 8px;
+}
+.progress-bar {
+  height: 100%;
+  background: var(--accent);
+  border-radius: 8px;
+  transition: width 0.3s;
+  width: 0%;
+}
+.progress-text {
+  font-size: 13px;
+  color: var(--text2);
+  font-family: 'SF Mono', monospace;
+}
+
+/* Info */
+.info-bar {
+  background: var(--surface);
+  border-radius: 8px;
+  padding: 10px 16px;
+  font-family: 'SF Mono', monospace;
+  font-size: 13px;
+  color: var(--text2);
+  display: none;
+}
+.info-bar.active { display: block; }
+
+/* Batch gallery */
+.gallery {
+  display: grid;
+  grid-template-columns: repeat(auto-fill, minmax(180px, 1fr));
+  gap: 12px;
+  margin-bottom: 16px;
+}
+.gallery-item {
+  background: var(--surface);
+  border-radius: 8px;
+  overflow: hidden;
+  cursor: pointer;
+  transition: transform 0.2s;
+}
+.gallery-item:hover { transform: scale(1.02); }
+.gallery-item img {
+  width: 100%;
+  aspect-ratio: 1;
+  object-fit: cover;
+}
+.gallery-item .name {
+  padding: 6px 10px;
+  font-size: 11px;
+  color: var(--text2);
+  white-space: nowrap;
+  overflow: hidden;
+  text-overflow: ellipsis;
+}
+
+/* File list */
+.file-list {
+  display: flex;
+  flex-wrap: wrap;
+  gap: 8px;
+  margin-top: 8px;
+}
+.file-chip {
+  background: var(--surface2);
+  border-radius: 6px;
+  padding: 4px 10px;
+  font-size: 12px;
+  color: var(--text2);
+  display: flex;
+  align-items: center;
+  gap: 6px;
+}
+.file-chip .remove {
+  cursor: pointer;
+  color: #f44;
+  font-weight: bold;
+}
+
+/* Spinner */
+.spinner {
+  display: inline-block;
+  width: 18px;
+  height: 18px;
+  border: 2px solid rgba(255,255,255,0.3);
+  border-top-color: #fff;
+  border-radius: 50%;
+  animation: spin 0.8s linear infinite;
+}
+@keyframes spin { to { transform: rotate(360deg); } }
+
+/* Toast */
+.toast {
+  position: fixed;
+  bottom: 24px;
+  right: 24px;
+  background: var(--surface2);
+  border: 1px solid var(--border);
+  border-radius: 8px;
+  padding: 12px 20px;
+  font-size: 14px;
+  z-index: 100;
+  display: none;
+  animation: slideUp 0.3s;
+}
+@keyframes slideUp {
+  from { transform: translateY(20px); opacity: 0; }
+  to { transform: translateY(0); opacity: 1; }
+}
+
+@media (max-width: 640px) {
+  .preview-area { grid-template-columns: 1fr; }
+  .control-row { flex-direction: column; align-items: stretch; }
+}
+</style>
+</head>
+<body>
+<div class="container">
+  <header>
+    <h1>Thera <span>MLX</span></h1>
+    <p>Arbitrary-scale super-resolution on Apple Silicon</p>
+  </header>
+
+  <div class="tabs">
+    <button class="tab active" onclick="switchTab('single')">Image</button>
+    <button class="tab" onclick="switchTab('batch')">Batch</button>
+    <button class="tab" onclick="switchTab('video')">Video</button>
+  </div>
+
+  <!-- ==================== SINGLE IMAGE ==================== -->
+  <div id="tab-single" class="tab-content active">
+    <div class="controls">
+      <div class="control-row">
+        <div class="control-group">
+          <label>Scale</label>
+          <div style="display:flex;align-items:center;gap:8px">
+            <input type="range" id="s-scale" min="1" max="8" step="0.5" value="2"
+              oninput="document.getElementById('s-scale-val').textContent=this.value+'x'">
+            <span class="scale-display" id="s-scale-val">2x</span>
+          </div>
+        </div>
+        <div class="control-group">
+          <label>Model</label>
+          <select id="s-model">
+            <option value="air">Air (fast)</option>
+            <option value="pro">Pro (quality)</option>
+          </select>
+        </div>
+        <label class="checkbox-label">
+          <input type="checkbox" id="s-ensemble"> Ensemble
+        </label>
+        <div class="control-group">
+          <label>Tiles</label>
+          <select id="s-tiles">
+            <option value="1">Off</option>
+            <option value="2">2x2</option>
+            <option value="3">3x3</option>
+            <option value="4">4x4</option>
+          </select>
+        </div>
+      </div>
+    </div>
+
+    <div class="dropzone" id="s-dropzone"
+      onclick="document.getElementById('s-file').click()">
+      <div class="dropzone-icon">🖼</div>
+      <div class="dropzone-text">
+        <strong>Drop image here</strong> or click to browse
+      </div>
+      <input type="file" id="s-file" accept="image/*" hidden
+        onchange="handleSingleFile(this.files[0])">
+    </div>
+
+    <div class="preview-area" id="s-preview" style="display:none">
+      <div class="preview-panel">
+        <span class="panel-label">Input</span>
+        <img id="s-input-img">
+      </div>
+      <div class="preview-panel">
+        <span class="panel-label">Output</span>
+        <img id="s-output-img">
+      </div>
+    </div>
+
+    <div class="compare-container" id="s-compare">
+      <span class="compare-label compare-label-before">Before</span>
+      <span class="compare-label compare-label-after">After</span>
+      <img id="s-compare-after" style="width:100%;display:block">
+      <div class="compare-overlay" id="s-compare-overlay">
+        <img id="s-compare-before">
+      </div>
+      <div class="compare-line" id="s-compare-line"></div>
+    </div>
+
+    <div class="progress-container" id="s-progress">
+      <div class="progress-bar-bg"><div class="progress-bar" id="s-pbar"></div></div>
+      <div class="progress-text" id="s-ptext">Processing...</div>
+    </div>
+
+    <div class="action-row">
+      <button class="btn btn-primary" id="s-btn" onclick="doSingleUpscale()" disabled>
+        Upscale
+      </button>
+      <button class="btn btn-secondary" id="s-compare-btn" style="display:none"
+        onclick="toggleCompare()">
+        Compare
+      </button>
+      <button class="btn btn-secondary" id="s-download" style="display:none"
+        onclick="downloadSingle()">
+        Download PNG
+      </button>
+    </div>
+
+    <div class="info-bar" id="s-info" style="margin-top:12px"></div>
+  </div>
+
+  <!-- ==================== BATCH ==================== -->
+  <div id="tab-batch" class="tab-content">
+    <div class="controls">
+      <div class="control-row">
+        <div class="control-group">
+          <label>Scale</label>
+          <div style="display:flex;align-items:center;gap:8px">
+            <input type="range" id="b-scale" min="1" max="8" step="0.5" value="2"
+              oninput="document.getElementById('b-scale-val').textContent=this.value+'x'">
+            <span class="scale-display" id="b-scale-val">2x</span>
+          </div>
+        </div>
+        <div class="control-group">
+          <label>Model</label>
+          <select id="b-model">
+            <option value="air">Air (fast)</option>
+            <option value="pro">Pro (quality)</option>
+          </select>
+        </div>
+        <label class="checkbox-label">
+          <input type="checkbox" id="b-ensemble"> Ensemble
+        </label>
+      </div>
+    </div>
+
+    <div class="dropzone" id="b-dropzone"
+      onclick="document.getElementById('b-file').click()">
+      <div class="dropzone-icon">📁</div>
+      <div class="dropzone-text">
+        <strong>Drop images here</strong> or click to browse (multiple)
+      </div>
+      <input type="file" id="b-file" accept="image/*" multiple hidden
+        onchange="handleBatchFiles(this.files)">
+      <div class="file-list" id="b-filelist"></div>
+    </div>
+
+    <div class="progress-container" id="b-progress">
+      <div class="progress-bar-bg"><div class="progress-bar" id="b-pbar"></div></div>
+      <div class="progress-text" id="b-ptext">Processing...</div>
+    </div>
+
+    <div class="action-row">
+      <button class="btn btn-primary" id="b-btn" onclick="doBatchUpscale()" disabled>
+        Upscale All
+      </button>
+      <button class="btn btn-secondary" id="b-download" style="display:none"
+        onclick="downloadBatch()">
+        Download ZIP
+      </button>
+    </div>
+
+    <div class="info-bar" id="b-info" style="margin-top:12px"></div>
+  </div>
+
+  <!-- ==================== VIDEO ==================== -->
+  <div id="tab-video" class="tab-content">
+    <div class="controls">
+      <div class="control-row">
+        <div class="control-group">
+          <label>Scale</label>
+          <div style="display:flex;align-items:center;gap:8px">
+            <input type="range" id="v-scale" min="1" max="4" step="0.5" value="2"
+              oninput="document.getElementById('v-scale-val').textContent=this.value+'x'">
+            <span class="scale-display" id="v-scale-val">2x</span>
+          </div>
+        </div>
+        <div class="control-group">
+          <label>Model</label>
+          <select id="v-model">
+            <option value="air">Air (fast)</option>
+            <option value="pro">Pro (quality)</option>
+          </select>
+        </div>
+      </div>
+    </div>
+
+    <div class="dropzone" id="v-dropzone"
+      onclick="document.getElementById('v-file').click()">
+      <div class="dropzone-icon">🎬</div>
+      <div class="dropzone-text">
+        <strong>Drop video here</strong> or click to browse
+      </div>
+      <input type="file" id="v-file" accept="video/*" hidden
+        onchange="handleVideoFile(this.files[0])">
+    </div>
+
+    <div class="preview-area" id="v-preview" style="display:none">
+      <div class="preview-panel">
+        <span class="panel-label">Input</span>
+        <video id="v-input-vid" controls muted></video>
+      </div>
+      <div class="preview-panel">
+        <span class="panel-label">Output</span>
+        <video id="v-output-vid" controls></video>
+      </div>
+    </div>
+
+    <div class="progress-container" id="v-progress">
+      <div class="progress-bar-bg"><div class="progress-bar" id="v-pbar"></div></div>
+      <div class="progress-text" id="v-ptext">Processing...</div>
+    </div>
+
+    <div class="action-row">
+      <button class="btn btn-primary" id="v-btn" onclick="doVideoUpscale()" disabled>
+        Upscale Video
+      </button>
+      <button class="btn btn-secondary" id="v-download" style="display:none"
+        onclick="downloadVideo()">
+        Download MP4
+      </button>
+    </div>
+
+    <div class="info-bar" id="v-info" style="margin-top:12px"></div>
+  </div>
+</div>
+
+<div class="toast" id="toast"></div>
+
+<script>
+// --- Tab switching ---
+function switchTab(tab) {
+  document.querySelectorAll('.tab').forEach((el, i) => {
+    const tabs = ['single','batch','video'];
+    el.classList.toggle('active', tabs[i] === tab);
+  });
+  document.querySelectorAll('.tab-content').forEach(el => el.classList.remove('active'));
+  document.getElementById('tab-' + tab).classList.add('active');
+}
+
+// --- Toast ---
+function toast(msg, duration=3000) {
+  const el = document.getElementById('toast');
+  el.textContent = msg;
+  el.style.display = 'block';
+  setTimeout(() => el.style.display = 'none', duration);
+}
+
+// --- Drag & drop ---
+document.querySelectorAll('.dropzone').forEach(dz => {
+  dz.addEventListener('dragover', e => { e.preventDefault(); dz.classList.add('dragover'); });
+  dz.addEventListener('dragleave', () => dz.classList.remove('dragover'));
+  dz.addEventListener('drop', e => {
+    e.preventDefault();
+    dz.classList.remove('dragover');
+    const input = dz.querySelector('input[type="file"]');
+    if (input.multiple) {
+      handleBatchFiles(e.dataTransfer.files);
+    } else if (input.accept.includes('video')) {
+      handleVideoFile(e.dataTransfer.files[0]);
+    } else {
+      handleSingleFile(e.dataTransfer.files[0]);
+    }
+  });
+});
+
+// ==================== SINGLE ====================
+let singleFile = null;
+let singleBlob = null;
+let singleInputUrl = null;
+let compareActive = false;
+
+function handleSingleFile(file) {
+  if (!file) return;
+  singleFile = file;
+  const dz = document.getElementById('s-dropzone');
+  dz.classList.add('has-file');
+  dz.querySelector('.dropzone-text').innerHTML = '<strong>' + file.name + '</strong>';
+
+  singleInputUrl = URL.createObjectURL(file);
+  document.getElementById('s-input-img').src = singleInputUrl;
+  document.getElementById('s-output-img').src = '';
+  document.getElementById('s-preview').style.display = 'grid';
+  document.getElementById('s-btn').disabled = false;
+  document.getElementById('s-download').style.display = 'none';
+  document.getElementById('s-compare-btn').style.display = 'none';
+  document.getElementById('s-info').classList.remove('active');
+  hideCompare();
+}
+
+async function doSingleUpscale() {
+  if (!singleFile) return;
+  const btn = document.getElementById('s-btn');
+  btn.disabled = true;
+  btn.innerHTML = '<span class="spinner"></span> Processing';
+  hideCompare();
+
+  const prog = document.getElementById('s-progress');
+  prog.classList.add('active');
+  document.getElementById('s-pbar').style.width = '60%';
+  const tiles = document.getElementById('s-tiles').value;
+  const tileLabel = tiles > 1 ? ` (${tiles}x${tiles} tiles)` : '';
+  document.getElementById('s-ptext').textContent = 'Upscaling' + tileLabel + '...';
+
+  const fd = new FormData();
+  fd.append('image', singleFile);
+  fd.append('scale', document.getElementById('s-scale').value);
+  fd.append('model', document.getElementById('s-model').value);
+  fd.append('ensemble', document.getElementById('s-ensemble').checked);
+  fd.append('tiles', tiles);
+
+  try {
+    const resp = await fetch('/api/upscale', { method: 'POST', body: fd });
+    if (!resp.ok) { const e = await resp.json(); throw new Error(e.error); }
+
+    const blob = await resp.blob();
+    singleBlob = blob;
+    const url = URL.createObjectURL(blob);
+    document.getElementById('s-output-img').src = url;
+
+    const info = JSON.parse(resp.headers.get('X-Info') || '{}');
+    document.getElementById('s-info').textContent =
+      `${info.src} \u2192 ${info.dst} (${info.scale}x) | ${info.model} | ${info.time}s`;
+    document.getElementById('s-info').classList.add('active');
+    document.getElementById('s-download').style.display = 'inline-flex';
+    document.getElementById('s-compare-btn').style.display = 'inline-flex';
+
+    document.getElementById('s-pbar').style.width = '100%';
+    toast('Upscale complete!');
+  } catch (e) {
+    toast('Error: ' + e.message, 5000);
+  } finally {
+    btn.disabled = false;
+    btn.innerHTML = 'Upscale';
+    prog.classList.remove('active');
+  }
+}
+
+function downloadSingle() {
+  if (!singleBlob) return;
+  const a = document.createElement('a');
+  a.href = URL.createObjectURL(singleBlob);
+  const name = singleFile.name.replace(/\.[^.]+$/, '') + '_thera.png';
+  a.download = name;
+  a.click();
+}
+
+// --- Compare split-view ---
+function toggleCompare() {
+  if (compareActive) {
+    hideCompare();
+  } else {
+    showCompare();
+  }
+}
+
+function showCompare() {
+  if (!singleBlob || !singleInputUrl) return;
+  compareActive = true;
+
+  const outputUrl = URL.createObjectURL(singleBlob);
+  const container = document.getElementById('s-compare');
+  const afterImg = document.getElementById('s-compare-after');
+  const beforeImg = document.getElementById('s-compare-before');
+
+  // Use the upscaled version to determine the display size,
+  // then scale the input up to match via CSS so it's a fair pixel comparison
+  afterImg.src = outputUrl;
+  beforeImg.src = singleInputUrl;
+
+  afterImg.onload = () => {
+    // Set before image width to match container
+    beforeImg.style.width = container.offsetWidth + 'px';
+    updateCompareSlider(0.5);
+  };
+
+  document.getElementById('s-preview').style.display = 'none';
+  container.classList.add('active');
+
+  document.getElementById('s-compare-btn').textContent = 'Side by Side';
+}
+
+function hideCompare() {
+  compareActive = false;
+  document.getElementById('s-compare').classList.remove('active');
+  if (singleFile) {
+    document.getElementById('s-preview').style.display = 'grid';
+  }
+  document.getElementById('s-compare-btn').textContent = 'Compare';
+}
+
+function updateCompareSlider(ratio) {
+  const container = document.getElementById('s-compare');
+  const overlay = document.getElementById('s-compare-overlay');
+  const line = document.getElementById('s-compare-line');
+  const w = container.offsetWidth;
+  const pos = Math.max(0, Math.min(w, w * ratio));
+
+  overlay.style.width = pos + 'px';
+  line.style.left = pos + 'px';
+}
+
+// Compare drag handling
+(function() {
+  const container = document.getElementById('s-compare');
+  let dragging = false;
+
+  function onMove(clientX) {
+    const rect = container.getBoundingClientRect();
+    const ratio = (clientX - rect.left) / rect.width;
+    updateCompareSlider(Math.max(0, Math.min(1, ratio)));
+  }
+
+  container.addEventListener('mousedown', e => { dragging = true; onMove(e.clientX); });
+  window.addEventListener('mousemove', e => { if (dragging) onMove(e.clientX); });
+  window.addEventListener('mouseup', () => { dragging = false; });
+
+  container.addEventListener('touchstart', e => { dragging = true; onMove(e.touches[0].clientX); }, {passive: true});
+  window.addEventListener('touchmove', e => { if (dragging) onMove(e.touches[0].clientX); }, {passive: true});
+  window.addEventListener('touchend', () => { dragging = false; });
+})();
+
+// ==================== BATCH ====================
+let batchFiles = [];
+let batchBlob = null;
+
+function handleBatchFiles(files) {
+  batchFiles = Array.from(files);
+  const dz = document.getElementById('b-dropzone');
+  dz.classList.add('has-file');
+
+  const list = document.getElementById('b-filelist');
+  list.innerHTML = batchFiles.map((f, i) =>
+    `<span class="file-chip">${f.name}
+      <span class="remove" onclick="event.stopPropagation();removeBatchFile(${i})">×</span>
+    </span>`
+  ).join('');
+
+  dz.querySelector('.dropzone-text').innerHTML =
+    `<strong>${batchFiles.length} image${batchFiles.length>1?'s':''}</strong> selected`;
+  document.getElementById('b-btn').disabled = false;
+}
+
+function removeBatchFile(idx) {
+  batchFiles.splice(idx, 1);
+  if (batchFiles.length === 0) {
+    const dz = document.getElementById('b-dropzone');
+    dz.classList.remove('has-file');
+    dz.querySelector('.dropzone-text').innerHTML =
+      '<strong>Drop images here</strong> or click to browse (multiple)';
+    document.getElementById('b-filelist').innerHTML = '';
+    document.getElementById('b-btn').disabled = true;
+  } else {
+    handleBatchFiles(batchFiles);
+  }
+}
+
+async function doBatchUpscale() {
+  if (!batchFiles.length) return;
+  const btn = document.getElementById('b-btn');
+  btn.disabled = true;
+  btn.innerHTML = '<span class="spinner"></span> Processing';
+
+  const prog = document.getElementById('b-progress');
+  prog.classList.add('active');
+  document.getElementById('b-pbar').style.width = '30%';
+  document.getElementById('b-ptext').textContent =
+    `Upscaling ${batchFiles.length} images...`;
+
+  const fd = new FormData();
+  batchFiles.forEach(f => fd.append('images', f));
+  fd.append('scale', document.getElementById('b-scale').value);
+  fd.append('model', document.getElementById('b-model').value);
+  fd.append('ensemble', document.getElementById('b-ensemble').checked);
+
+  try {
+    const resp = await fetch('/api/batch', { method: 'POST', body: fd });
+    if (!resp.ok) { const e = await resp.json(); throw new Error(e.error); }
+
+    batchBlob = await resp.blob();
+    const info = JSON.parse(resp.headers.get('X-Info') || '{}');
+    document.getElementById('b-info').textContent =
+      `${info.count} images | ${info.scale}x | ${info.model} | ${info.time}s`;
+    document.getElementById('b-info').classList.add('active');
+    document.getElementById('b-download').style.display = 'inline-flex';
+    document.getElementById('b-pbar').style.width = '100%';
+
+    toast(`Batch complete — ${info.count} images`);
+  } catch (e) {
+    toast('Error: ' + e.message, 5000);
+  } finally {
+    btn.disabled = false;
+    btn.innerHTML = 'Upscale All';
+    prog.classList.remove('active');
+  }
+}
+
+function downloadBatch() {
+  if (!batchBlob) return;
+  const a = document.createElement('a');
+  a.href = URL.createObjectURL(batchBlob);
+  a.download = 'thera_batch.zip';
+  a.click();
+}
+
+// ==================== VIDEO ====================
+let videoFile = null;
+let videoJobId = null;
+
+function handleVideoFile(file) {
+  if (!file) return;
+  videoFile = file;
+  const dz = document.getElementById('v-dropzone');
+  dz.classList.add('has-file');
+  dz.querySelector('.dropzone-text').innerHTML = '<strong>' + file.name + '</strong>';
+
+  const url = URL.createObjectURL(file);
+  document.getElementById('v-input-vid').src = url;
+  document.getElementById('v-output-vid').src = '';
+  document.getElementById('v-preview').style.display = 'grid';
+  document.getElementById('v-btn').disabled = false;
+  document.getElementById('v-download').style.display = 'none';
+  document.getElementById('v-info').classList.remove('active');
+}
+
+async function doVideoUpscale() {
+  if (!videoFile) return;
+  const btn = document.getElementById('v-btn');
+  btn.disabled = true;
+  btn.innerHTML = '<span class="spinner"></span> Processing';
+
+  const prog = document.getElementById('v-progress');
+  prog.classList.add('active');
+  document.getElementById('v-pbar').style.width = '2%';
+  document.getElementById('v-ptext').textContent = 'Uploading...';
+
+  const fd = new FormData();
+  fd.append('video', videoFile);
+  fd.append('scale', document.getElementById('v-scale').value);
+  fd.append('model', document.getElementById('v-model').value);
+
+  try {
+    const resp = await fetch('/api/video/start', { method: 'POST', body: fd });
+    if (!resp.ok) { const e = await resp.json(); throw new Error(e.error); }
+    const data = await resp.json();
+    videoJobId = data.job_id;
+    pollVideoProgress();
+  } catch (e) {
+    toast('Error: ' + e.message, 5000);
+    btn.disabled = false;
+    btn.innerHTML = 'Upscale Video';
+    prog.classList.remove('active');
+  }
+}
+
+async function pollVideoProgress() {
+  if (!videoJobId) return;
+
+  try {
+    const resp = await fetch('/api/video/progress/' + videoJobId);
+    const job = await resp.json();
+
+    const pbar = document.getElementById('v-pbar');
+    const ptext = document.getElementById('v-ptext');
+
+    if (job.status === 'extracting') {
+      pbar.style.width = '5%';
+      ptext.textContent = 'Extracting frames...';
+    } else if (job.status === 'upscaling') {
+      const pct = 5 + 85 * (job.current_frame / Math.max(job.total_frames, 1));
+      pbar.style.width = pct + '%';
+      ptext.textContent =
+        `Frame ${job.current_frame}/${job.total_frames} | ` +
+        `${job.fps} fps | ETA ${job.eta}s`;
+    } else if (job.status === 'encoding') {
+      pbar.style.width = '92%';
+      ptext.textContent = 'Encoding video...';
+    } else if (job.status === 'done') {
+      pbar.style.width = '100%';
+      ptext.textContent = 'Complete!';
+      document.getElementById('v-output-vid').src =
+        '/api/video/download/' + videoJobId;
+      document.getElementById('v-info').textContent =
+        `${job.src} → ${job.dst} (${job.scale}x) | ` +
+        `${job.total_frames} frames | ${job.time}s | ${job.model}`;
+      document.getElementById('v-info').classList.add('active');
+      document.getElementById('v-download').style.display = 'inline-flex';
+
+      document.getElementById('v-btn').disabled = false;
+      document.getElementById('v-btn').innerHTML = 'Upscale Video';
+      toast('Video upscale complete!');
+      return;
+    } else if (job.status === 'error') {
+      toast('Error: ' + job.error, 5000);
+      document.getElementById('v-btn').disabled = false;
+      document.getElementById('v-btn').innerHTML = 'Upscale Video';
+      document.getElementById('v-progress').classList.remove('active');
+      return;
+    }
+
+    setTimeout(pollVideoProgress, 1000);
+  } catch (e) {
+    setTimeout(pollVideoProgress, 2000);
+  }
+}
+
+function downloadVideo() {
+  if (!videoJobId) return;
+  const a = document.createElement('a');
+  a.href = '/api/video/download/' + videoJobId;
+  a.download = 'thera_upscaled.mp4';
+  a.click();
+}
+</script>
+</body>
+</html>
+"""
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+if __name__ == "__main__":
+    import logging
+    logging.getLogger("werkzeug").setLevel(logging.ERROR)
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--port", type=int, default=5005)
+    parser.add_argument("--host", type=str, default="127.0.0.1")
+    args = parser.parse_args()
+
+    print(f"Thera MLX  →  http://{args.host}:{args.port}")
+    app.run(host=args.host, port=args.port, debug=False, threaded=True)

upscale.py

@@ -0,0 +1,186 @@
+#!/usr/bin/env python3
+"""Core upscale logic for Thera MLX."""
+
+import time
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+from PIL import Image
+
+from model import Thera
+
+WEIGHTS_DIR = Path(__file__).parent / "weights"
+
+
+def load_weights(model, weights_path):
+    """Load converted weights into the MLX model."""
+    weights_path = str(weights_path)
+    if weights_path.endswith('.safetensors'):
+        from safetensors.numpy import load_file
+        raw = load_file(weights_path)
+        weights = {k: mx.array(v) for k, v in raw.items()}
+    elif weights_path.endswith('.npz'):
+        raw = np.load(weights_path)
+        weights = {k: mx.array(raw[k]) for k, v in raw.items()}
+    else:
+        raise ValueError(f"Unknown weight format: {weights_path}")
+
+    weight_list = list(weights.items())
+    model.load_weights(weight_list)
+    return model
+
+
+def get_weights_path(model_size):
+    """Resolve weights path for a model variant."""
+    return WEIGHTS_DIR / f"weights-{model_size}.safetensors"
+
+
+def upscale_tiled(model, source_np, target_h, target_w, tiles, ensemble=False):
+    """Upscale an image using NxN tiles to reduce peak RAM.
+
+    Splits the source image into tiles with overlap, upscales each tile
+    individually, then blends them back together using linear feathering.
+
+    Args:
+        model: Loaded Thera model.
+        source_np: numpy array (H, W, 3) float32 in [0, 1].
+        target_h: Target height.
+        target_w: Target width.
+        tiles: Number of tiles per axis (2, 3, or 4).
+        ensemble: Use geometric self-ensemble.
+
+    Returns:
+        numpy uint8 array (target_h, target_w, 3).
+    """
+    h, w = source_np.shape[:2]
+    scale_h = target_h / h
+    scale_w = target_w / w
+
+    # Overlap in source pixels (10% of tile size, minimum 8px)
+    tile_h = h / tiles
+    tile_w = w / tiles
+    overlap_h = max(8, int(tile_h * 0.1))
+    overlap_w = max(8, int(tile_w * 0.1))
+
+    # Build output canvas (float32 for blending)
+    output = np.zeros((target_h, target_w, 3), dtype=np.float32)
+    weight_map = np.zeros((target_h, target_w, 1), dtype=np.float32)
+
+    total_tiles = tiles * tiles
+    done = 0
+
+    for row in range(tiles):
+        for col in range(tiles):
+            # Source tile bounds with overlap
+            sy0 = round(row * h / tiles) - (overlap_h if row > 0 else 0)
+            sy1 = round((row + 1) * h / tiles) + (overlap_h if row < tiles - 1 else 0)
+            sx0 = round(col * w / tiles) - (overlap_w if col > 0 else 0)
+            sx1 = round((col + 1) * w / tiles) + (overlap_w if col < tiles - 1 else 0)
+
+            sy0 = max(0, sy0)
+            sy1 = min(h, sy1)
+            sx0 = max(0, sx0)
+            sx1 = min(w, sx1)
+
+            tile_src = source_np[sy0:sy1, sx0:sx1]
+            th = round((sy1 - sy0) * scale_h)
+            tw = round((sx1 - sx0) * scale_w)
+
+            # Upscale tile
+            result = model.upscale(mx.array(tile_src), th, tw, ensemble=ensemble)
+            mx.eval(result)
+            tile_out = np.array(result).astype(np.float32) / 255.0
+
+            # Target tile bounds
+            ty0 = round(sy0 * scale_h)
+            tx0 = round(sx0 * scale_w)
+            ty1 = ty0 + tile_out.shape[0]
+            tx1 = tx0 + tile_out.shape[1]
+
+            # Clamp to output bounds
+            ty1 = min(ty1, target_h)
+            tx1 = min(tx1, target_w)
+            tile_out = tile_out[:ty1 - ty0, :tx1 - tx0]
+
+            # Linear feather weight for blending overlaps
+            fh, fw = tile_out.shape[:2]
+            wy = np.ones(fh, dtype=np.float32)
+            wx = np.ones(fw, dtype=np.float32)
+
+            # Feather top/bottom edges in overlap regions
+            ovl_top = round(overlap_h * scale_h) if row > 0 else 0
+            ovl_bot = round(overlap_h * scale_h) if row < tiles - 1 else 0
+            ovl_left = round(overlap_w * scale_w) if col > 0 else 0
+            ovl_right = round(overlap_w * scale_w) if col < tiles - 1 else 0
+
+            if ovl_top > 0:
+                ramp = np.linspace(0, 1, min(ovl_top, fh), dtype=np.float32)
+                wy[:len(ramp)] = ramp
+            if ovl_bot > 0:
+                ramp = np.linspace(1, 0, min(ovl_bot, fh), dtype=np.float32)
+                wy[-len(ramp):] = np.minimum(wy[-len(ramp):], ramp)
+            if ovl_left > 0:
+                ramp = np.linspace(0, 1, min(ovl_left, fw), dtype=np.float32)
+                wx[:len(ramp)] = ramp
+            if ovl_right > 0:
+                ramp = np.linspace(1, 0, min(ovl_right, fw), dtype=np.float32)
+                wx[-len(ramp):] = np.minimum(wx[-len(ramp):], ramp)
+
+            w2d = wy[:, None] * wx[None, :]  # (fh, fw)
+            w3d = w2d[:, :, None]  # (fh, fw, 1)
+
+            output[ty0:ty1, tx0:tx1] += tile_out * w3d
+            weight_map[ty0:ty1, tx0:tx1] += w3d
+
+            done += 1
+            print(f"  tile {done}/{total_tiles}")
+
+    # Normalize by weight
+    weight_map = np.maximum(weight_map, 1e-8)
+    output = (output / weight_map * 255 + 0.5).clip(0, 255).astype(np.uint8)
+    return output
+
+
+def upscale_file(input_path, output_path, scale=None, size=None,
+                 model_size='air', weights_path=None, ensemble=False,
+                 tiles=None):
+    """Upscale a single image file."""
+    img = Image.open(input_path).convert('RGB')
+    source = np.asarray(img, dtype=np.float32) / 255.0
+    h, w = source.shape[:2]
+
+    if scale is not None:
+        target_h = round(h * scale)
+        target_w = round(w * scale)
+    elif size is not None:
+        target_h, target_w = size
+    else:
+        raise ValueError("Must specify either scale or size")
+
+    scale_actual = target_h / h
+    if weights_path is None:
+        weights_path = get_weights_path(model_size)
+
+    model = Thera(size=model_size)
+    model = load_weights(model, weights_path)
+    mx.eval(model.parameters())
+
+    t0 = time.perf_counter()
+
+    if tiles and tiles > 1:
+        print(f"Tiled upscale: {tiles}x{tiles} ({tiles*tiles} tiles)")
+        result_np = upscale_tiled(model, source, target_h, target_w,
+                                  tiles, ensemble=ensemble)
+        Image.fromarray(result_np).save(output_path)
+    else:
+        result = model.upscale(mx.array(source), target_h, target_w, ensemble=ensemble)
+        mx.eval(result)
+        Image.fromarray(np.array(result)).save(output_path)
+
+    elapsed = time.perf_counter() - t0
+
+    suffix = " (ensemble)" if ensemble else ""
+    tile_info = f" [{tiles}x{tiles} tiles]" if tiles and tiles > 1 else ""
+    print(f"[{model_size}]{suffix}{tile_info} {w}x{h} -> {target_w}x{target_h} ({scale_actual:.4g}x)  {elapsed:.1f}s  ->  {output_path}")