$ErrorActionPreference = 'Stop'

$envPath = Join-Path $PSScriptRoot '..\..\..\.env'
Get-Content -LiteralPath $envPath | ForEach-Object {
    if ($_ -match '^\s*([^#=\s]+)\s*=\s*(.*)\s*$') {
        [Environment]::SetEnvironmentVariable($matches[1], $matches[2].Trim('"').Trim("'"), 'Process')
    }
}
if (-not $env:RUNPOD_API_KEY) { throw 'RUNPOD_API_KEY is not set in .env.' }

$probe = @'
import json
import torch
from safetensors import safe_open
import comfy_kitchen as ck
from comfy_kitchen.tensor import TensorCoreNVFP4Layout

path = "/runpod-volume/ComfyUI/models/diffusion_models/minimax_h3_ref2va_pruned_nvfp4.safetensors"
prefix = "blocks.0.attn.qkv_proj"
with safe_open(path, framework="pt", device="cuda") as checkpoint:
    sidecar = checkpoint.get_tensor(prefix + ".comfy_quant")
    assert json.loads(bytes(sidecar.cpu().tolist())) == {"format": "nvfp4"}
    weight = checkpoint.get_tensor(prefix + ".weight").contiguous()
    weight_scale = checkpoint.get_tensor(prefix + ".weight_scale").view(torch.float8_e4m3fn).contiguous()
    weight_scale_2 = checkpoint.get_tensor(prefix + ".weight_scale_2").to(torch.float32).contiguous()

x = torch.randn((1, weight.shape[1] * 2), device="cuda", dtype=torch.bfloat16)
packed_x, x_params = TensorCoreNVFP4Layout.quantize(x)
y = ck.scaled_mm_nvfp4(
    packed_x, weight,
    tensor_scale_a=x_params.scale,
    tensor_scale_b=weight_scale_2,
    block_scale_a=x_params.block_scale,
    block_scale_b=weight_scale,
    out_dtype=torch.bfloat16,
)
torch.cuda.synchronize()
print(json.dumps({
    "input_shape": list(x.shape),
    "packed_weight_shape": list(weight.shape),
    "weight_scale_shape": list(weight_scale.shape),
    "weight_scale_2_shape": list(weight_scale_2.shape),
    "output_shape": list(y.shape),
    "output_dtype": str(y.dtype),
    "finite": bool(torch.isfinite(y).all()),
    "gpu": torch.cuda.get_device_name(),
}, sort_keys=True))
'@
$encodedProbe = [Convert]::ToBase64String([Text.Encoding]::UTF8.GetBytes($probe))
$setup = 'python -m pip install --no-cache-dir comfy-kitchen==0.2.28 "safetensors>=0.5.0"; python -c "import base64;exec(base64.b64decode(\"' + $encodedProbe + '\"))"'
$payload = @{
    name = 'h3-nvfp4-linear-probe'
    image = 'runpod/pytorch:1.1.0-rc.154-cu1300-torch291-ubuntu2404'
    cloud = 'SECURE'
    gpu = @{ id = 'NVIDIA RTX PRO 6000 Blackwell Server Edition'; count = 1 }
    disk = 20
    dataCenterIds = @('EUR-IS-1')
    mounts = @{ network = @(@{ volumeId = 'g8r2uqufz5'; path = '/runpod-volume' }) }
    args = 'bash -lc ''set -eu; ' + $setup + ''''
}

$headers = @{ Authorization = "Bearer $env:RUNPOD_API_KEY" }
Invoke-RestMethod -Method Post -Uri 'https://api.runpod.io/v2/pods' -Headers $headers -ContentType 'application/json' -Body ($payload | ConvertTo-Json -Depth 8)