FoundationVision/LlamaGen

1# Modified from:2#   gpt-fast: https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py3#   DiT:      https://github.com/facebookresearch/DiT/blob/main/models.py4import torch5import torch.nn as nn6from torch.nn import functional as F7import torch._dynamo.config8import torch._inductor.config9import copy10# torch._inductor.config.coordinate_descent_tuning = True11# torch._inductor.config.triton.unique_kernel_names = True12# torch._inductor.config.fx_graph_cache = True # Experimental feature to reduce compilation times, will be on by default in future13 14 15### from https://huggingface.co/transformers/v3.2.0/_modules/transformers/generation_utils.html16def top_k_top_p_filtering(17    logits,18    top_k: int = 0,19    top_p: float = 1.0,20    filter_value: float = -float("Inf"),21    min_tokens_to_keep: int = 1,22):23    """Filter a distribution of logits using top-k and/or nucleus (top-p) filtering24    Args:25        logits: logits distribution shape (batch size, vocabulary size)26        if top_k > 0: keep only top k tokens with highest probability (top-k filtering).27        if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).28            Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)29        Make sure we keep at least min_tokens_to_keep per batch example in the output30    From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf31731    """32    if top_k > 0:33        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))  # Safety check34        # Remove all tokens with a probability less than the last token of the top-k35        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]36        logits[indices_to_remove] = filter_value37 38    if top_p < 1.0:39        sorted_logits, sorted_indices = torch.sort(logits, descending=True)40        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)41 42        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)43        sorted_indices_to_remove = cumulative_probs > top_p44        if min_tokens_to_keep > 1:45            # Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)46            sorted_indices_to_remove[..., :min_tokens_to_keep] = 047        # Shift the indices to the right to keep also the first token above the threshold48        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()49        sorted_indices_to_remove[..., 0] = 050 51        # scatter sorted tensors to original indexing52        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)53        logits[indices_to_remove] = filter_value54    return logits55 56 57def sample(logits, temperature: float=1.0, top_k: int=0, top_p: float=1.0, sample_logits=True):        58    logits = logits[:, -1, :] / max(temperature, 1e-5)59    if top_k > 0 or top_p < 1.0:60        logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)61    probs = F.softmax(logits, dim=-1)62    if sample_logits:63        idx = torch.multinomial(probs, num_samples=1)64    else:65        _, idx = torch.topk(probs, k=1, dim=-1)66    return idx, probs67 68 69def logits_to_probs(logits, temperature: float = 1.0, top_p: float=1.0, top_k: int = None, **kwargs):70    logits = logits / max(temperature, 1e-5)71    if top_k > 0 or top_p < 1.0:72        logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p)73    probs = torch.nn.functional.softmax(logits, dim=-1)74    return probs75 76 77def prefill(model, cond_idx: torch.Tensor, input_pos: torch.Tensor, cfg_scale: float, **sampling_kwargs):78    if cfg_scale > 1.0:79        logits, _ = model(None, cond_idx, input_pos)80        logits_combined = logits81        cond_logits, uncond_logits = torch.split(logits_combined, len(logits_combined) // 2, dim=0)82        logits = uncond_logits + (cond_logits - uncond_logits) * cfg_scale83    else:84        logits, _ = model(None, cond_idx, input_pos)85 86    return sample(logits, **sampling_kwargs)[0]87 88 89def decode_one_token(model, x: torch.Tensor, input_pos: torch.Tensor, cfg_scale: float, cfg_flag: bool, **sampling_kwargs):90    assert input_pos.shape[-1] == 191    if cfg_scale > 1.0:92        x_combined = torch.cat([x, x])93        logits, _ = model(x_combined, cond_idx=None, input_pos=input_pos)94        logits_combined = logits95        cond_logits, uncond_logits = torch.split(logits_combined, len(logits_combined) // 2, dim=0) 96        if cfg_flag:97            logits = uncond_logits + (cond_logits - uncond_logits) * cfg_scale98        else:99            logits = cond_logits100    else:101        logits, _ = model(x, cond_idx=None, input_pos=input_pos)102    return sample(logits, **sampling_kwargs)103 104 105def decode_n_tokens(106    model, cur_token: torch.Tensor, input_pos: torch.Tensor, num_new_tokens: int, 107    cfg_scale: float, cfg_interval: int,108    **sampling_kwargs):109    new_tokens, new_probs = [], []110    cfg_flag = True111    for i in range(num_new_tokens):112        with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True): # Actually better for Inductor to codegen attention here113            if cfg_interval > -1 and i > cfg_interval:114                cfg_flag = False115            next_token, next_prob = decode_one_token(116                model, cur_token, input_pos, cfg_scale, cfg_flag, **sampling_kwargs117            )118            input_pos += 1119            new_tokens.append(next_token.clone())120            new_probs.append(next_prob.clone())121            cur_token = next_token.view(-1, 1)122    123    return new_tokens, new_probs124 125 126@torch.no_grad()127def generate(model, cond, max_new_tokens, emb_masks=None, cfg_scale=1.0, cfg_interval=-1, **sampling_kwargs):128    if model.model_type == 'c2i':129        if cfg_scale > 1.0:130            cond_null = torch.ones_like(cond) * model.num_classes131            cond_combined = torch.cat([cond, cond_null])132        else:133            cond_combined = cond134        T = 1135    elif model.model_type == 't2i':136        if cfg_scale > 1.0:137            cond_null = torch.zeros_like(cond) + model.cls_embedding.uncond_embedding138            cond_combined = torch.cat([cond, cond_null])139        else:140            cond_combined = cond141        T = cond.shape[1]      142    else:143        raise Exception("please check model type")144 145    T_new = T + max_new_tokens146    max_seq_length = T_new147    max_batch_size = cond.shape[0]148 149    device = cond.device150    with torch.device(device):151        max_batch_size_cfg = max_batch_size * 2 if cfg_scale > 1.0 else max_batch_size152        model.setup_caches(max_batch_size=max_batch_size_cfg, max_seq_length=max_seq_length, dtype=model.tok_embeddings.weight.dtype)153    154    if emb_masks is not None:155        assert emb_masks.shape[0] == max_batch_size156        assert emb_masks.shape[-1] == T157        if cfg_scale > 1.0:158            model.causal_mask[:, :, :T] = model.causal_mask[:, :, :T] * torch.cat([emb_masks, emb_masks]).unsqueeze(1)159        else:160            model.causal_mask[:, :, :T] = model.causal_mask[:, :, :T] * emb_masks.unsqueeze(1)161 162        eye_matrix = torch.eye(model.causal_mask.size(1), model.causal_mask.size(2), device=device)163        model.causal_mask[:] = model.causal_mask * (1 - eye_matrix) + eye_matrix164    165    # create an empty tensor of the expected final shape and fill in the current tokens166    seq = torch.empty((max_batch_size, T_new), dtype=torch.int, device=device)167 168    input_pos = torch.arange(0, T, device=device)169    next_token = prefill(model, cond_combined, input_pos, cfg_scale, **sampling_kwargs)170    seq[:, T:T+1] = next_token171 172    input_pos = torch.tensor([T], device=device, dtype=torch.int)173    generated_tokens, _ = decode_n_tokens(model, next_token, input_pos, max_new_tokens-1, cfg_scale, cfg_interval, **sampling_kwargs)174    seq[:, T+1:] = torch.cat(generated_tokens, dim=1)175 176    return seq[:, T:]177