Function infer

app.py:53–90

(cfg_scale, top_k, top_p, temperature, class_label, seed)

def infer(cfg_scale, top_k, top_p, temperature, class_label, seed):
    llm.llm_engine.model_executor.driver_worker.model_runner.model.sampler = Sampler(cfg_scale)
    args.cfg_scale = cfg_scale
    n = 4
    latent_size = image_size // args.downsample_size
    # Labels to condition the model with (feel free to change):
    class_labels = [class_label for _ in range(n)]
    qzshape = [len(class_labels), args.codebook_embed_dim, latent_size, latent_size]

    prompt_token_ids = [[cind] for cind in class_labels]
    if cfg_scale > 1.0:
        prompt_token_ids.extend([[args.num_classes] for _ in range(len(prompt_token_ids))])

    # Create a sampling params object.
    sampling_params = SamplingParams(
        temperature=temperature, top_p=top_p, top_k=top_k, 
        max_tokens=latent_size ** 2)
    
    t1 = time.time()
    torch.manual_seed(seed)
    outputs = llm.generate(
        prompt_token_ids=prompt_token_ids,
        sampling_params=sampling_params,
        use_tqdm=False)
    sampling_time = time.time() - t1
    print(f"gpt sampling takes about {sampling_time:.2f} seconds.")    

    index_sample = torch.tensor([output.outputs[0].token_ids for output in outputs], device=device)
    if cfg_scale > 1.0:
        index_sample = index_sample[:len(class_labels)]
    t2 = time.time()
    samples = vq_model.decode_code(index_sample, qzshape) # output value is between [-1, 1]
    decoder_time = time.time() - t2
    print(f"decoder takes about {decoder_time:.2f} seconds.")
    # Convert to PIL.Image format:
    samples = samples.mul(127.5).add_(128.0).clamp_(0, 255).permute(0, 2, 3, 1).to("cpu", torch.uint8).numpy()
    samples = [Image.fromarray(sample) for sample in samples]
    return samples


parser = argparse.ArgumentParser()