export_pid.py

import os
import sys
import torch
import torch.nn as nn
import argparse

# Add PiD repository to python path
sys.path.append(os.path.join(os.path.dirname(__file__), "PiD_repo"))

# MOCK text encoder loading to bypass Gemma-2-2B HF gated license requirements
import pid._src.models.pixeldit_model as pixeldit_model

class DummyTokenizer:
    def encode(self, text):
        return []

# Replace text encoder loader with a dummy
pixeldit_model._load_text_encoder = lambda name, device="cuda": (DummyTokenizer(), None)

# Replace text encoder function to return zeros
def dummy_encode_text_raw(self, captions):
    # SDXL config uses model_max_length=300 and caption_channels=2304
    B = len(captions)
    dummy_embs = torch.zeros(B, 300, 2304, device="cuda", dtype=torch.float32)
    dummy_masks = torch.ones(B, 300, device="cuda", dtype=torch.float32)
    return dummy_embs, dummy_masks

pixeldit_model.PixelDiTModel._encode_text_raw = dummy_encode_text_raw

# MOCK VAE loading to avoid loading './checkpoints/sdxl_vae.safetensors'
import pid._src.tokenizers.sdxl_vae as sdxl_vae
sdxl_vae._sdxl_vae = lambda pretrained_path=None, device="cpu", s3_credential_path=None: sdxl_vae.SDAutoEncoder(sdxl_vae.SDXL_VAE_PARAMS).to(device)

from pid._src.utils.model_loader import load_model_from_checkpoint

class PidNetWrapper(nn.Module):
    def __init__(self, net, precomputed_y):
        super().__init__()
        self.net = net
        # Register the empty prompt context as a buffer (constant in ONNX)
        self.register_buffer("y", precomputed_y)

    def forward(self, x, t, lq_latent, degrade_sigma):
        # x: [1, 3, 4096, 4096]
        # t: [1] (scaled timestep)
        # lq_latent: [1, 4, 128, 128]
        # degrade_sigma: [1]
        return self.net(
            x=x,
            t=t,
            y=self.y,
            lq_latent=lq_latent,
            degrade_sigma=degrade_sigma
        )

def export_pid(checkpoint_path, output_path):
    print("Loading PiD model with mocked VAE and text encoder...")
    
    model, config = load_model_from_checkpoint(
        experiment_name="PiD_res2kto4k_sr4x_official_sdxl_distill_4step",
        checkpoint_path=checkpoint_path,
        config_file="PiD_repo/pid/_src/configs/pid/config.py",
        enable_fsdp=False,
        strict=False,
        load_ema_to_reg=True
    )
    model.eval()
    
    print("Generating zero prompt text embedding...")
    # Get zero prompt embedding
    with torch.no_grad():
        caption_embs, _ = model._encode_text_raw([""])
        caption_embs = caption_embs.cuda()
        
    print(f"Text embedding shape: {caption_embs.shape}") # Should be [1, 300, 2304]
    
    # Extract the net and wrap it
    net = model.net
    net.eval()
    
    wrapper = PidNetWrapper(net, caption_embs)
    wrapper.cuda()
    wrapper.eval()
    
    # Dummy inputs for ONNX export
    dummy_x = torch.zeros(1, 3, 4096, 4096, dtype=torch.float32, device="cuda")
    dummy_t = torch.tensor([999.0], dtype=torch.float32, device="cuda") # scaled t
    dummy_lq_latent = torch.zeros(1, 4, 128, 128, dtype=torch.float32, device="cuda")
    dummy_degrade_sigma = torch.tensor([0.0], dtype=torch.float32, device="cuda")
    
    print(f"Exporting PiD to ONNX: {output_path}")
    
    torch.onnx.export(
        wrapper,
        (dummy_x, dummy_t, dummy_lq_latent, dummy_degrade_sigma),
        output_path,
        export_params=True,
        opset_version=18,
        do_constant_folding=True,
        input_names=["x", "t", "lq_latent", "degrade_sigma"],
        output_names=["v_pred"],
        dynamic_axes=None
    )
    print("PiD model ONNX export complete!")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Export PiD Model to ONNX")
    parser.add_argument("--checkpoint", type=str, required=True,
                        help="Path to consolidated checkpoint (.pth) file")
    parser.add_argument("--output", type=str, default="pid_upscaler.onnx",
                        help="Path to output ONNX file")
    args = parser.parse_args()
    
    export_pid(args.checkpoint, args.output)