#!/usr/bin/env python3
"""Single-prompt SD3-M / SD3.5-M generation with the matched released adapter.

Install the image package from https://github.com/TWWinde/ReSeed first.
No remote Python model code is loaded by this script.
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--weights", type=Path, required=True, help="downloaded single-seed snapshot")
    parser.add_argument("--model", choices=["sd3m", "sd35m"], required=True)
    parser.add_argument("--base", type=Path, required=True, help="downloaded upstream teacher snapshot")
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--seed", type=int, default=0, help="generation noise seed, not training seed")
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--without-reseed", action="store_true")
    args = parser.parse_args()
    if args.output.exists():
        raise FileExistsError("refusing to overwrite an existing image")
    import torch
    from diffusers import StableDiffusion3Pipeline, SD3Transformer2DModel
    from reseed.models.sd3_residual_bias import load_sd3_residual_bias_adapter

    entry = json.loads((args.weights/"release_manifest.json").read_text())["models"][args.model]
    transformer = SD3Transformer2DModel.from_pretrained(
        str(args.weights/entry["student_path"]/"student"), torch_dtype=torch.bfloat16)
    pipe = StableDiffusion3Pipeline.from_pretrained(str(args.base), transformer=transformer,
                                                   torch_dtype=torch.bfloat16).to("cuda")
    if abs(float(pipe.scheduler.config.shift)-3.) > 1e-7 or pipe.scheduler.config.get("use_dynamic_shifting", False):
        raise ValueError("this release uses static FlowMatch shift=3, not dynamic shifting")
    for component in pipe.components.values():
        if isinstance(component, torch.nn.Module):
            component.eval().requires_grad_(False)
    adapter = None
    if not args.without_reseed:
        adapter, _ = load_sd3_residual_bias_adapter(pipe.transformer, args.weights/entry["adapter_path"])
        adapter.to("cuda").eval().requires_grad_(False)
        adapter.enable()
    def callback(_pipe, index, _timestep, kwargs):
        if adapter is not None:
            adapter.disable()  # Step zero just finished; never activate on steps 1–3.
        return kwargs
    shape = (1, int(transformer.config.in_channels), 1024//pipe.vae_scale_factor, 1024//pipe.vae_scale_factor)
    noise = torch.randn(shape, generator=torch.Generator(device="cpu").manual_seed(args.seed), dtype=torch.float32)
    try:
        with torch.inference_mode():
            image = pipe(prompt=[args.prompt], negative_prompt=[""], height=1024, width=1024,
                         sigmas=[1., .75, .5, .25], guidance_scale=1., max_sequence_length=256,
                         latents=noise.to(device="cuda", dtype=torch.bfloat16),
                         callback_on_step_end=callback if adapter is not None else None,
                         callback_on_step_end_tensor_inputs=["latents"]).images[0]
        args.output.parent.mkdir(parents=True, exist_ok=True)
        image.save(args.output)
    finally:
        if adapter is not None:
            adapter.disable(); adapter.close()


if __name__ == "__main__":
    main()