import cv2
import torch
from torch.nn import functional as F
import os

CONFIG = {
    'img0': '/path/to/image0.png',
    'img1': '/path/to/image1.png', 
    'model_dir': 'train_log_wms',
    'out': 'out_mid.png',
}

try:
    from train_log.RIFE_HDv3 import Model      # prefer v3
except Exception:
    try:
        from model.RIFE_HDv2 import Model
    except Exception:
        try:
            from model.RIFE_HD import Model
        except Exception:
            from model.RIFE import Model


def pad32(t):
    n, c, h, w = t.shape
    ph = ((h - 1) // 32 + 1) * 32
    pw = ((w - 1) // 32 + 1) * 32
    pad = (0, pw - w, 0, ph - h)
    return F.pad(t, pad), pad, (h, w)

def main():
    cfg = CONFIG
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    model = Model()
    model.load_model(cfg['model_dir'], -1)
    model.eval(); model.device()

    I0 = cv2.imread(cfg['img0'], cv2.IMREAD_UNCHANGED)
    I1 = cv2.imread(cfg['img1'], cv2.IMREAD_UNCHANGED)
    if I0.ndim == 2: I0 = cv2.cvtColor(I0, cv2.COLOR_GRAY2RGB)
    if I1.ndim == 2: I1 = cv2.cvtColor(I1, cv2.COLOR_GRAY2RGB)
    I0 = torch.from_numpy(I0[:,:,::-1].transpose(2,0,1)).float().unsqueeze(0)/255.0
    I1 = torch.from_numpy(I1[:,:,::-1].transpose(2,0,1)).float().unsqueeze(0)/255.0

    I0, pad, orig = pad32(I0); I1, _, _ = pad32(I1)

    with torch.no_grad():
        mid = model.inference(I0.to(device), I1.to(device))
    mid = (mid[0].cpu() * 255).byte().numpy().transpose(1,2,0)
    h, w = orig
    mid = mid[:h, :w, ::-1]
    os.makedirs(os.path.dirname(cfg['out']) or '.', exist_ok=True)
    cv2.imwrite(cfg['out'], mid)
    print(f"Saved {cfg['out']}")


if __name__ == '__main__':
    main()