Inference Code Bug, collapse when demo_num > 1.

[Kevin-lkw]

Hi! I find a code bug in inference.py:
It seems that context length is 16 frame_actions (16*347) , so we can set args.demo_num > 1 and set args.frams to 16-args.demo_num, and generate should work well.
However, when I set args.demo_num to some number > 1, the image collapse quickly.
possible reason:
I find that this is because When training(and inference), condition is in [image, action,image,action ... ] format. However, in inference code, condition (variable image_input) contains only images, and I guess this is OOD and caused collapse.
my solution: add actions between images, and it works well!

[Kevin-lkw]

I write a code to inference for longer than 16 frames, base on original code. I test with args.demo_num = 15 and args.fames=1 and it works. I past my code here, and it might have some personal setting but the whole idea is clear. I hope you can add code for inference longer than 16 frames!

def lvm_generate(args, model, input_path, output_path, file):

    ### previous code ###
    ...
    ### 3. load video frames 
    # num stand for args.demo_num

    # load serveral condition frames , shape[num,3,224,384]
    
    with torch.no_grad(), torch.autocast(device_type="cuda", dtype=torch.float16):
        img_index = model.tokenizer.tokenize_images(frames)
        #img_index shape [num ,14, 24]
    img_index = rearrange(img_index, '(b t) h w -> b t (h w)', b=1)
    # img_index shape 1, num ,336



    ### CENTAL PART OF MODIFICATION###   
    # action_list [N,11]
    # image_action_input [1, (num-1) * 347 + 336]
    image_action_input = []
    for i in range(args.demo_num):
        image_action_input.append(img_index[0][i])
        action_index = start_frame + i
        if i < args.demo_num - 1:
            image_action_input.append(torch.tensor(action_list[action_index]).to(device))

    # import ipdb; ipdb.set_trace()
    image_action_input = torch.cat(image_action_input, dim=0).unsqueeze(0)
    # shape [1,(num-1) * 347 + 336]
    ### END CENTAL PART OF MODIFICATION ###
    

    all_generated_tokens = []
    total = len(action_list)
    step = args.frames
    for i in range(end_frame, total, step):
        # currently predicting [i, i + step)
        cur_step = total - i if i + step > total else step
        action_all = action_list[i - 1: i + cur_step - 1] # as action i is performed after image i(in my dataset)
        action_all = torch.tensor(action_all).unsqueeze(1).to(device)
        
        start_t = time.time()
        with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.float16):
            if args.accelerate_algo == 'naive':
                outputs = model.transformer.naive_generate( input_ids=image_action_input, max_new_tokens=TOKEN_PER_PIX*cur_step, action_all=action_all, top_k=args.top_k, top_p=args.top_p)
            elif args.accelerate_algo == 'image_diagd':
                outputs = model.transformer.img_diagd_generate(input_ids=image_action_input, max_new_tokens=TOKEN_PER_PIX*cur_step, action_all=action_all,windowsize = args.window_size, top_k=args.top_k, top_p=args.top_p)
            else:
                raise ValueError(f"Unknown accelerate algorithm {args.accelerate_algo}")
        end_t = time.time()
        print(f"use {end_t - start_t} seconds to generate {cur_step} frames")
        all_generated_tokens.extend(outputs.tolist()[0])


#### THIS PART IS IMPORTANT TOO ####
        for j in range(cur_step):
            image_action_input = torch.cat([image_action_input, torch.tensor(action_all[j]).to(device)], dim=1)
            image_action_input = torch.cat([image_action_input, outputs[:,j*TOKEN_PER_PIX:(j+1)*TOKEN_PER_PIX]], dim=1)
        image_action_input = image_action_input[:,-((args.demo_num-1)*TOKEN_PER_IMAGE + TOKEN_PER_PIX):]
#### END IMPORTANT PART ####


    new_length = len(all_generated_tokens)
    time_costed = end_t - start_t 
    token_per_sec = new_length / time_costed
    frame_per_sec = token_per_sec / TOKEN_PER_PIX
    print(f"{new_length} token generated; cost {time_costed:.3f} second; {token_per_sec:.3f} token/sec {frame_per_sec:.3f} fps")
    token2video(all_generated_tokens, model.tokenizer, output_mp4_path, args.fps, device)  
    # return for evaluation 
    return_item = {
        "time_costed": time_costed,
        "token_num": new_length,
    }
    import ipdb; ipdb.set_trace()
    return return_item

[sysuyy]

Thank you for your correction! We design inference.py to reproduce only evaluation results, so we did not try setting args.frames to larger. We will test your code and merge into the our repo as soon as possible. Thanks again for your code !!!