# /// script
# requires-python = ">=3.11"
# dependencies = ["coreai-core==1.0.0b3", "coreai-torch==0.4.3", "spandrel==0.4.1", "numpy", "pillow"]
# [tool.uv]
# index-url = "https://pypi.org/simple"
# prerelease = "allow"
# index-strategy = "unsafe-best-match"
# ///
"""Export an ESRGAN-family upscaler (RRDBNet / SRVGGNetCompact, via spandrel) as a fixed-tile
Core AI asset with the same contract swift-upscale's CoreAIUpscaler runs for SwinIR-L:
input "x" [0,1] RGB, output "output" [0,1] RGB, clamped.
    uv run export_esrgan.py <ckpt.pth> <out.aimodel> <tile> <float16|float32> <author> <licence>
"""
import sys, time
from pathlib import Path
import torch
from spandrel import ModelLoader
from coreai.runtime import AIModelAssetMetadata
from coreai_torch import TorchConverter, get_decomp_table

ckpt, out, tile, dtype_name, author, licence = sys.argv[1], Path(sys.argv[2]), int(sys.argv[3]), sys.argv[4], sys.argv[5], sys.argv[6]
dtype = {"float16": torch.float16, "float32": torch.float32}[dtype_name]
desc = ModelLoader().load_from_file(ckpt)
print("arch", type(desc.model).__name__, "scale", desc.scale, flush=True)

class Wrapped(torch.nn.Module):
    """RGB in, RGB out. A greyscale model (one input channel — MangaJPEGLQ is built for
    black-and-white manga) sees Rec. 601 luma, and its single output channel is copied to
    all three, so the Swift side feeds and reads RGB either way."""
    def __init__(self, m, in_ch, out_ch):
        super().__init__(); self.m = m; self.in_ch = in_ch; self.out_ch = out_ch
        self.register_buffer("luma", torch.tensor([0.299, 0.587, 0.114]).view(1, 3, 1, 1))
    def forward(self, x):
        if self.in_ch == 1:
            x = (x * self.luma.to(x.dtype)).sum(dim=1, keepdim=True)
        y = self.m(x)
        if self.out_ch == 1:
            y = y.repeat(1, 3, 1, 1)
        return y.clamp(0, 1)

print("channels", desc.input_channels, "->", desc.output_channels, flush=True)
model = Wrapped(desc.model.eval(), desc.input_channels, desc.output_channels).to(dtype)
example = {"x": torch.rand(1, 3, tile, tile).to(dtype)}
t0 = time.time()
with torch.no_grad(), torch.autocast(device_type="cpu", dtype=dtype, enabled=dtype != torch.float32):
    exported = torch.export.export(model, args=(), kwargs=example)
exported = exported.run_decompositions(get_decomp_table())
program = TorchConverter(mode=TorchConverter.Mode.RELEASE).add_exported_program(
    exported_program=exported, input_names=["x"], output_names=["output"]).to_coreai()
md = AIModelAssetMetadata()
md.author = author
md.license = licence
md.model_description = (f"{Path(ckpt).stem}: x{desc.scale} super-resolution, {tile}x{tile} -> "
                        f"{tile*desc.scale}x{tile*desc.scale} tile, {dtype_name}, [0,1] RGB in/out.")
program.save_asset(out, metadata=md)
print(f"saved {out} in {time.time()-t0:.1f} s", flush=True)