"""One next-frame edit using the exported adapter and pinned DiffSynth-Studio."""
import argparse,pathlib
import torch
from PIL import Image
from diffsynth.pipelines.qwen_image_21 import QwenImage21Pipeline,ModelConfig

ap=argparse.ArgumentParser()
ap.add_argument('--base',required=True,help='Directory of the original BF16 Qwen-Image-2.1 snapshot')
ap.add_argument('--adapter',required=True)
ap.add_argument('--anchor',required=True)
ap.add_argument('--previous',required=True)
ap.add_argument('--prompt',required=True)
ap.add_argument('--output',default='next-frame.png')
ap.add_argument('--seed',type=int,default=1234)
args=ap.parse_args();base=pathlib.Path(args.base)
configs=[]
for name,pattern in [('transformer','diffusion_pytorch_model*.safetensors'),('text_encoder','model*.safetensors'),('vae','diffusion_pytorch_model*.safetensors')]:
    paths=sorted(str(p) for p in (base/name).glob(pattern));assert paths
    configs.append(ModelConfig(path=paths,skip_download=True))
pipe=QwenImage21Pipeline.from_pretrained(torch_dtype=torch.bfloat16,device='cuda',model_configs=configs,processor_config=ModelConfig(path=str(base/'processor'),skip_download=True))
pipe.load_lora(pipe.dit,args.adapter,alpha=1.0)  # 128 tensors should be fused.
previous=Image.open(args.previous).convert('RGBA')
size=(previous.width//32*32,previous.height//32*32)
refs=[Image.open(args.anchor).convert('RGBA').resize(size),previous.resize(size)]
result=pipe(prompt=args.prompt,edit_image=refs,width=size[0],height=size[1],seed=args.seed,num_inference_steps=40)
result.save(args.output)