# Fast CLIP-guided diffusion. Same math as the repo's CLIP-guided diffusion baseline
# (OpenAI 256/512 uncond model, p_sample + cond_fn nudging each step toward CLIP text),
# restructured for modern GPUs: one UNet forward per step instead of two (p_mean_variance
# and cond_fn share the same deterministic model output), bf16 conv torso, branch-free
# fixed-shape cutouts with randomness passed in as tensors, and one full timestep
# (UNet + guidance grad + DDPM update) as a single functional graph (torch.func.vjp/grad)
# that torch.compile's or AOTI-compiles via aokit (warm start = load-only, no JIT).

import argparse
import sys
import time

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from torchvision.transforms import functional as TF

sys.path.append('guided-diffusion')

_argv = sys.argv
sys.argv = [sys.argv[0]]  # generate_fast parses CLI args at import
from generate_fast import clip, aoti_build_or_load, convert_clip_visual  # noqa: E402
sys.argv = _argv

from guided_diffusion.script_util import create_model_and_diffusion, model_and_diffusion_defaults  # noqa: E402

torch.backends.cudnn.benchmark = True
torch.set_float32_matmul_precision('high')
device = torch.device('cuda:0')


def parse():
    p = argparse.ArgumentParser(description='Fast CLIP-guided diffusion')
    p.add_argument('-p', '--prompt', type=str, default='a photograph of a lighthouse in a storm', dest='prompt')
    p.add_argument('--steps', type=str, default='250', help='timestep_respacing (e.g. 250, 100)')
    p.add_argument('--image_size', type=int, default=512, help='256 or 512 (match the model)')
    p.add_argument('-m', '--clip_model', type=str, default='ViT-B/16', dest='clip_model')
    p.add_argument('--clip_guidance_scale', type=float, default=1000.)
    p.add_argument('--tv_scale', type=float, default=150.)
    p.add_argument('--range_scale', type=float, default=50.)
    p.add_argument('-cuts', '--num_cuts', type=int, default=16, dest='cutn')
    p.add_argument('--cut_pow', type=float, default=1.)
    p.add_argument('-sd', '--seed', type=int, default=0, dest='seed')
    p.add_argument('-o', '--output', type=str, default='diffusion_fast.png', dest='output')
    p.add_argument('-se', '--save_every', type=int, default=50, dest='save_every')
    p.add_argument('--model', type=str, default=None, help='checkpoint path (auto by image_size)')
    p.add_argument('--backend', choices=['baseline', 'eager', 'compile', 'aoti'], default='aoti', dest='backend')
    p.add_argument('--aoti_cache', type=str, default='aoti_cache', dest='aoti_cache')
    p.add_argument('--fp32', action='store_true', help='disable bf16 for the UNet/CLIP')
    p.add_argument('--eval', action='store_true', help='print fp32 CLIP score of the final image')
    return p.parse_args()


MODEL_CFG = {
    'attention_resolutions': '32, 16, 8', 'class_cond': False, 'diffusion_steps': 1000,
    'rescale_timesteps': True, 'learn_sigma': True, 'noise_schedule': 'linear',
    'num_channels': 256, 'num_head_channels': 64, 'num_res_blocks': 2,
    'resblock_updown': True, 'use_fp16': False, 'use_scale_shift_norm': True,
}


def load_diffusion(image_size, steps, ckpt):
    cfg = model_and_diffusion_defaults()
    cfg.update(MODEL_CFG)
    cfg.update({'image_size': image_size, 'timestep_respacing': str(steps)})
    model, diffusion = create_model_and_diffusion(**cfg)
    model.load_state_dict(torch.load(ckpt, map_location='cpu'))
    model.requires_grad_(False).eval().to(device)
    return model, diffusion


def convert_unet_bf16(model):
    # bf16 analog of convert_to_fp16: conv torso in bf16, norms/linears/out stay fp32
    for m in [model.input_blocks, model.middle_block, model.output_blocks]:
        for l in m.modules():
            if isinstance(l, (nn.Conv1d, nn.Conv2d, nn.Conv3d)):
                l.weight.data = l.weight.data.bfloat16()
                if l.bias is not None:
                    l.bias.data = l.bias.data.bfloat16()
    model.dtype = torch.bfloat16
    return model


def tv_loss(inp):
    inp = F.pad(inp, (0, 1, 0, 1), 'replicate')
    x = inp[..., :-1, 1:] - inp[..., :-1, :-1]
    y = inp[..., 1:, :-1] - inp[..., :-1, :-1]
    return (x ** 2 + y ** 2).mean([1, 2, 3])


def range_loss(inp):
    return (inp - inp.clamp(-1, 1)).pow(2).mean([1, 2, 3])


class FusedStep(nn.Module):
    """One full CLIP-guided DDPM timestep as a pure graph. The respaced schedule lives
    on the host: each step's scalars (UNet timestep value, alpha/posterior coefs,
    log-variance bounds, t!=0 noise mask) and the guidance scales enter as tensor
    inputs -> one AOTI package per (image_size, cutn, dtype) serves any steps/scales."""

    def __init__(self, model, clip_visual, cutn, cut_size, image_size, heavy_dtype):
        super().__init__()
        self.model = convert_unet_bf16(model) if heavy_dtype == torch.bfloat16 else model
        self.visual = convert_clip_visual(clip_visual.float(), heavy_dtype)
        self.cutn, self.cut_size, self.image_size = cutn, cut_size, image_size
        self.heavy_dtype = heavy_dtype
        self.register_buffer('norm_mean', torch.tensor([0.48145466, 0.4578275, 0.40821073]).view(1, 3, 1, 1))
        self.register_buffer('norm_std', torch.tensor([0.26862954, 0.26130258, 0.27577711]).view(1, 3, 1, 1))
        self.register_buffer('lin', torch.arange(cut_size).float())

    def make_cutouts(self, img, size, offx, offy):
        # random resized crop (branch-free, fixed cut_size output) via grid_sample;
        # size/offx/offy sampled to match MakeCutouts(cut_pow) distribution
        cs, n = self.cut_size, self.cutn
        H = W = self.image_size
        px = offx[:, None] + (self.lin[None, :] + 0.5) * size[:, None] / cs - 0.5   # [n, cs]
        py = offy[:, None] + (self.lin[None, :] + 0.5) * size[:, None] / cs - 0.5
        gx = (px + 0.5) * 2 / W - 1
        gy = (py + 0.5) * 2 / H - 1
        grid = torch.stack([gx[:, None, :].expand(n, cs, cs),
                            gy[:, :, None].expand(n, cs, cs)], dim=-1)            # [n, cs, cs, 2]
        batch = img.expand(n, -1, -1, -1)
        return F.grid_sample(batch, grid, mode='bilinear', padding_mode='border', align_corners=False)

    def _model_output(self, x, model_ts):
        out = self.model(x, model_ts)
        eps, var_values = torch.split(out, 3, dim=1)
        return eps.float(), var_values.float()

    def producer(self, x, model_ts, sqrt_recip, sqrt_recipm1, fac):
        eps, var_values = self._model_output(x, model_ts)
        pred_xstart = sqrt_recip * x - sqrt_recipm1 * eps
        x_in = pred_xstart * fac + x * (1 - fac)
        return (x_in, pred_xstart), var_values

    def guide(self, x_in, embeds, weights, size, offx, offy, cgs, tvs):
        x01 = x_in.add(1).div(2)
        cuts = self.make_cutouts(x01, size, offx, offy)
        cuts = (cuts - self.norm_mean) / self.norm_std
        emb = self.visual(cuts.to(self.heavy_dtype)).float()
        img_n = F.normalize(emb, dim=-1)                    # [cutn, D]
        txt_n = F.normalize(embeds, dim=-1)                 # [P, D]
        dists = (img_n.unsqueeze(1) - txt_n.unsqueeze(0)).norm(dim=-1).div(2).arcsin().pow(2).mul(2)  # [cutn, P]
        losses = (dists * weights).sum(1).mean(0)
        return losses * cgs + tv_loss(x_in).sum() * tvs

    def forward(self, x, model_ts, sqrt_recip, sqrt_recipm1, fac, pmc1, pmc2,
                min_log, max_log, nonzero, embeds, weights, cgs, tvs, rs,
                noise, size, offx, offy):
        (x_in, pred_xstart), vjp_fn, var_values = torch.func.vjp(
            lambda xx: self.producer(xx, model_ts, sqrt_recip, sqrt_recipm1, fac), x, has_aux=True)
        g_xin = torch.func.grad(
            lambda xi: self.guide(xi, embeds, weights, size, offx, offy, cgs, tvs))(x_in)
        # range_loss(pred_xstart) cotangent in closed form (d/dp (p-clamp(p))^2 = 2(p-clamp(p)));
        # note the baseline's range term was dead code (grad of a non-descendant wrt x_in = 0),
        # this applies the intended stabilizer; rs=0 reproduces the old graph exactly
        g_pred = rs * 2 * (pred_xstart - pred_xstart.clamp(-1, 1)) / pred_xstart[0].numel()
        guidance = -vjp_fn((g_xin, g_pred))[0]              # -grad of guidance loss wrt x

        # LEARNED_RANGE variance from the same forward
        frac = (var_values + 1) / 2
        model_log_variance = frac * max_log + (1 - frac) * min_log
        model_variance = torch.exp(model_log_variance)

        model_mean = pmc1 * pred_xstart + pmc2 * x
        new_mean = model_mean + model_variance * guidance
        sample = new_mean + nonzero * torch.exp(0.5 * model_log_variance) * noise
        return sample, pred_xstart


def make_sched(diffusion, dev):
    # host-side respaced schedule tables; per-step scalars are fed to the graph as inputs
    d = diffusion

    def T(a):
        return torch.tensor(np.asarray(a), dtype=torch.float32, device=dev)

    model_ts = T(d.timestep_map)
    if d.rescale_timesteps:
        model_ts = model_ts * (1000.0 / d.original_num_steps)
    nz = torch.ones(d.num_timesteps, device=dev)
    nz[0] = 0.
    return {
        'model_ts': model_ts,
        'sqrt_recip': T(d.sqrt_recip_alphas_cumprod),
        'sqrt_recipm1': T(d.sqrt_recipm1_alphas_cumprod),
        'fac': T(d.sqrt_one_minus_alphas_cumprod),
        'pmc1': T(d.posterior_mean_coef1),
        'pmc2': T(d.posterior_mean_coef2),
        'min_log': T(d.posterior_log_variance_clipped),
        'max_log': T(np.log(d.betas)),
        'nz': nz,
    }


def sched_at(s, i):
    return (s['model_ts'][i:i + 1], s['sqrt_recip'][i], s['sqrt_recipm1'][i], s['fac'][i],
            s['pmc1'][i], s['pmc2'][i], s['min_log'][i], s['max_log'][i], s['nz'][i])


def sample_step_rands(gen, cutn, cut_size, image_size, cut_pow, shape):
    max_size, min_size = image_size, min(image_size, cut_size)
    u = torch.rand(cutn, generator=gen, device=device)
    size = (u.pow(cut_pow) * (max_size - min_size) + min_size).floor()
    offx = (torch.rand(cutn, generator=gen, device=device) * (image_size - size + 1)).floor()
    offy = (torch.rand(cutn, generator=gen, device=device) * (image_size - size + 1)).floor()
    noise = torch.randn(shape, generator=gen, device=device)
    return noise, size, offx, offy


def encode_text(perceptor, prompt):
    embeds, weights = [], []
    for part in prompt.split('|'):
        part = part.strip()
        if not part:
            continue
        vals = part.rsplit(':', 1)
        txt = vals[0]
        w = float(vals[1]) if len(vals) == 2 else 1.0
        with torch.no_grad():
            embeds.append(perceptor.encode_text(clip.tokenize(txt).to(device)).float()[0])
        weights.append(w)
    embeds = torch.stack(embeds)
    weights = torch.tensor(weights, device=device)
    weights = weights / weights.sum().abs()
    return embeds, weights


# ---------- baseline (reference): literal 2-forward fp32 loop via the diffusion object ----------

def run_baseline(args, model, diffusion, perceptor, cut_size, embeds, weights):
    mean = torch.tensor([0.48145466, 0.4578275, 0.40821073], device=device).view(1, 3, 1, 1)
    std = torch.tensor([0.26862954, 0.26130258, 0.27577711], device=device).view(1, 3, 1, 1)
    isz, cutn, cp = args.image_size, args.cutn, args.cut_pow

    def make_cutouts(img):
        cuts = []
        for _ in range(cutn):
            size = int(torch.rand([])**cp * (isz - min(isz, cut_size)) + min(isz, cut_size))
            ox = torch.randint(0, isz - size + 1, ())
            oy = torch.randint(0, isz - size + 1, ())
            c = img[:, :, oy:oy + size, ox:ox + size]
            cuts.append(F.adaptive_avg_pool2d(c, cut_size))
        return torch.cat(cuts)

    total = diffusion.num_timesteps
    cur_t = [0]

    def cond_fn(x, t, y=None):
        with torch.enable_grad():
            x = x.detach().requires_grad_()
            my_t = torch.ones([x.shape[0]], device=device, dtype=torch.long) * cur_t[0]
            out = diffusion.p_mean_variance(model, x, my_t, clip_denoised=False, model_kwargs={})
            fac = diffusion.sqrt_one_minus_alphas_cumprod[cur_t[0]]
            x_in = out['pred_xstart'] * fac + x * (1 - fac)
            clip_in = (make_cutouts(x_in.add(1).div(2)) - mean) / std
            image_embeds = perceptor.encode_image(clip_in).float()
            d = (F.normalize(image_embeds.unsqueeze(1), dim=-1) - F.normalize(embeds.unsqueeze(0), dim=-1))
            dists = d.norm(dim=-1).div(2).arcsin().pow(2).mul(2)
            losses = (dists * weights).sum(1).mean(0)
            x_in_grad = torch.autograd.grad(losses.sum() * args.clip_guidance_scale, x_in)[0]
            loss = tv_loss(x_in).sum() * args.tv_scale + range_loss(out['pred_xstart']).sum() * args.range_scale
            x_in_grad = x_in_grad + torch.autograd.grad(loss, x_in)[0]
            return -torch.autograd.grad(x_in, x, x_in_grad)[0]

    cur_t[0] = diffusion.num_timesteps - 1
    t0 = None
    img = None
    for j, sample in enumerate(diffusion.p_sample_loop_progressive(
            model, (1, 3, isz, isz), clip_denoised=False, model_kwargs={},
            cond_fn=cond_fn, progress=False, randomize_class=False)):
        if j == 5:
            torch.cuda.synchronize(); t0 = time.time()
        cur_t[0] -= 1
        img = sample['pred_xstart']
        if j % args.save_every == 0:
            TF.to_pil_image(img.add(1).div(2).clamp(0, 1)[0].cpu()).save(args.output)
    torch.cuda.synchronize()
    n = total - 5
    print(f'baseline: {n} steps in {time.time()-t0:.2f}s = {n/(time.time()-t0):.2f} step/s (fp32, {isz}px)')
    final = img.add(1).div(2).clamp(0, 1)
    TF.to_pil_image(final[0].cpu()).save(args.output)
    return final


# ---------- fast path (eager / compile / aoti) ----------

def build_engine(args, step_mod, embeds, weights):
    # steps/scales are NOT in the key: one package per (clip, ckpt, size, cutn, dtype)
    isz = args.image_size
    shape = (1, 3, isz, isz)
    gen = torch.Generator(device=device).manual_seed(0)
    ex = sample_step_rands(gen, args.cutn, step_mod.cut_size, isz, args.cut_pow, shape)

    # every example input must be a DISTINCT tensor object: make_fx binds repeated
    # tensors to one placeholder, silently miswiring the runtime inputs
    def s(v):
        return torch.full((), float(v), device=device)

    ex_sched = (torch.full((1,), 999., device=device),
                s(1.1), s(1.2), s(0.9), s(0.5), s(0.6), s(-9.), s(-8.), s(1.))
    ex_args = (torch.randn(shape, device=device), *ex_sched, embeds, weights,
               s(1000.), s(150.), s(50.), *ex)
    aoti = aoti_build_or_load('diffstep3-dyn', lambda *a: step_mod(*a), ex_args,
                              (args.clip_model, args.model, isz, args.cutn, str(step_mod.heavy_dtype)),
                              args.aoti_cache)
    return lambda *a: aoti(*a)


def run_fast(args, step_mod, diffusion, embeds, weights):
    isz = args.image_size
    total = diffusion.num_timesteps
    shape = (1, 3, isz, isz)

    if args.backend == 'compile':
        step_mod.forward = torch.compile(step_mod.forward, fullgraph=True)
        step_call = step_mod
    elif args.backend == 'aoti':
        step_call = build_engine(args, step_mod, embeds, weights)
    else:
        step_call = step_mod

    sched = make_sched(diffusion, device)
    cgs = torch.tensor(args.clip_guidance_scale, device=device)
    tvs = torch.tensor(args.tv_scale, device=device)
    rs = torch.tensor(args.range_scale, device=device)

    gen = torch.Generator(device=device).manual_seed(args.seed)
    x = torch.randn(shape, generator=gen, device=device)
    t0 = None
    img = None
    for j, i in enumerate(reversed(range(total))):
        noise, size, offx, offy = sample_step_rands(gen, args.cutn, step_mod.cut_size, isz, args.cut_pow, shape)
        x, pred_xstart = step_call(x, *sched_at(sched, i), embeds, weights, cgs, tvs, rs,
                                   noise, size, offx, offy)
        x = x.contiguous()
        img = pred_xstart
        if j == 5:
            torch.cuda.synchronize(); t0 = time.time()
        if j % args.save_every == 0:
            TF.to_pil_image(img.add(1).div(2).clamp(0, 1)[0].float().cpu()).save(args.output)
    torch.cuda.synchronize()
    n = total - 5
    dt = time.time() - t0
    print(f'{args.backend}: {n} steps in {dt:.2f}s = {n/dt:.2f} step/s '
          f'({"fp32" if args.fp32 else "bf16"}, {isz}px, cutn {args.cutn})')
    final = img.add(1).div(2).clamp(0, 1).float()
    TF.to_pil_image(final[0].cpu()).save(args.output)
    return final


@torch.no_grad()
def clip_score(clip_model, image01, prompt):
    # clean fp32 CLIP cosine of the final image vs the prompt (fresh model: the main
    # perceptor's visual was mutated to bf16 in place by convert_clip_visual)
    perceptor = clip.load(clip_model, jit=False)[0].eval().float().to(device)
    txt = perceptor.encode_text(clip.tokenize(prompt.split('|')[0].rsplit(':', 1)[0]).to(device)).float()
    img = F.interpolate(image01, 224, mode='bicubic', align_corners=False)
    mean = torch.tensor([0.48145466, 0.4578275, 0.40821073], device=device).view(1, 3, 1, 1)
    std = torch.tensor([0.26862954, 0.26130258, 0.27577711], device=device).view(1, 3, 1, 1)
    emb = perceptor.encode_image((img - mean) / std).float()
    return F.cosine_similarity(F.normalize(emb, dim=-1), F.normalize(txt, dim=-1)).item()


def main():
    args = parse()
    if args.model is None:
        args.model = ('checkpoints/256x256_diffusion_uncond.pt' if args.image_size == 256
                      else 'checkpoints/512x512_diffusion_uncond_finetune_008100.pt')
    heavy_dtype = torch.float32 if args.fp32 else torch.bfloat16
    torch.manual_seed(args.seed)
    print(f'CLIP-guided diffusion | {args.image_size}px | steps {args.steps} | backend {args.backend} | '
          f'cgs {args.clip_guidance_scale} tv {args.tv_scale} range {args.range_scale}')

    model, diffusion = load_diffusion(args.image_size, args.steps, args.model)
    perceptor = clip.load(args.clip_model, jit=False)[0].eval().requires_grad_(False).to(device)
    cut_size = perceptor.visual.input_resolution
    embeds, weights = encode_text(perceptor, args.prompt)
    print(f'prompt: {args.prompt}')

    if args.backend == 'baseline':
        final = run_baseline(args, model, diffusion, perceptor, cut_size, embeds, weights)
    else:
        step_mod = FusedStep(model, perceptor.visual, args.cutn, cut_size,
                             args.image_size, heavy_dtype).to(device).eval()
        final = run_fast(args, step_mod, diffusion, embeds, weights)

    print(f'saved {args.output}')
    if args.eval:
        print(f'CLIP score (fp32, vs prompt): {clip_score(args.clip_model, final, args.prompt):.4f}')


if __name__ == '__main__':
    main()