"""Encode ONE fixed prompt with cosmos T5 and save {t5_text_embeddings,t5_text_mask} pkl,
so training can set text_encoder_path='' (T5 never loaded, ~22GB freed) and use this cache."""
import argparse, pickle, torch
from cosmos_predict2.auxiliary.text_encoder import CosmosT5TextEncoder

def main(a):
    enc = CosmosT5TextEncoder(device='cuda', cache_dir=a.t5_path)
    enc.to('cuda')
    emb, mask = enc.encode_prompts([a.prompt], max_length=512, return_mask=True)
    # ONLY tensor keys — a4d_vae calls .to() on every value in this dict, so no strings allowed.
    payload = {'t5_text_embeddings': emb.to(torch.bfloat16).cpu(),
               't5_text_mask': mask.to(torch.bfloat16).cpu()}
    with open(a.out, 'wb') as f:
        pickle.dump(payload, f)
    print('emb', tuple(emb.shape), 'mask', tuple(mask.shape), 'nonzero-tok', int(mask.sum().item()))
    print('wrote', a.out)
    print('BUILD_T5_OK')

if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--prompt', default='pick up the cup and place it on the saucer')
    p.add_argument('--t5_path', required=True)
    p.add_argument('--out', default='custom/vae/default_text_pp70.pkl')
    main(p.parse_args())