Spaces:

THUDM
/

CogVideo

Running

CogVideo / model.py

hysts

Add files

8aac645 almost 2 years ago

No virus

51.9 kB

	# This code is adapted from https://github.com/THUDM/CogView2/blob/4e55cce981eb94b9c8c1f19ba9f632fd3ee42ba8/cogview2_text2image.py

	from __future__ import annotations

	import argparse
	import functools
	import logging
	import pathlib
	import sys
	import tempfile
	import time
	from typing import Any

	import gradio as gr
	import imageio.v2 as iio
	import numpy as np
	import torch
	from icetk import IceTokenizer
	from SwissArmyTransformer import get_args
	from SwissArmyTransformer.arguments import set_random_seed
	from SwissArmyTransformer.generation.sampling_strategies import BaseStrategy
	from SwissArmyTransformer.resources import auto_create

	app_dir = pathlib.Path(__file__).parent
	submodule_dir = app_dir / 'CogVideo'
	sys.path.insert(0, submodule_dir.as_posix())

	from coglm_strategy import CoglmStrategy
	from models.cogvideo_cache_model import CogVideoCacheModel
	from sr_pipeline import DirectSuperResolution

	formatter = logging.Formatter(
	'[%(asctime)s] %(name)s %(levelname)s: %(message)s',
	datefmt='%Y-%m-%d %H:%M:%S')
	stream_handler = logging.StreamHandler(stream=sys.stdout)
	stream_handler.setLevel(logging.INFO)
	stream_handler.setFormatter(formatter)
	logger = logging.getLogger(__name__)
	logger.setLevel(logging.INFO)
	logger.propagate = False
	logger.addHandler(stream_handler)

	ICETK_MODEL_DIR = app_dir / 'icetk_models'


	def get_masks_and_position_ids_stage1(data, textlen, framelen):
	# Extract batch size and sequence length.
	tokens = data
	seq_length = len(data[0])
	# Attention mask (lower triangular).
	attention_mask = torch.ones((1, textlen + framelen, textlen + framelen),
	device=data.device)
	attention_mask[:, :textlen, textlen:] = 0
	attention_mask[:, textlen:, textlen:].tril_()
	attention_mask.unsqueeze_(1)
	# Unaligned version
	position_ids = torch.zeros(seq_length,
	dtype=torch.long,
	device=data.device)
	torch.arange(textlen,
	out=position_ids[:textlen],
	dtype=torch.long,
	device=data.device)
	torch.arange(512,
	512 + seq_length - textlen,
	out=position_ids[textlen:],
	dtype=torch.long,
	device=data.device)
	position_ids = position_ids.unsqueeze(0)

	return tokens, attention_mask, position_ids


	def get_masks_and_position_ids_stage2(data, textlen, framelen):
	# Extract batch size and sequence length.
	tokens = data
	seq_length = len(data[0])

	# Attention mask (lower triangular).
	attention_mask = torch.ones((1, textlen + framelen, textlen + framelen),
	device=data.device)
	attention_mask[:, :textlen, textlen:] = 0
	attention_mask[:, textlen:, textlen:].tril_()
	attention_mask.unsqueeze_(1)

	# Unaligned version
	position_ids = torch.zeros(seq_length,
	dtype=torch.long,
	device=data.device)
	torch.arange(textlen,
	out=position_ids[:textlen],
	dtype=torch.long,
	device=data.device)
	frame_num = (seq_length - textlen) // framelen
	assert frame_num == 5
	torch.arange(512,
	512 + framelen,
	out=position_ids[textlen:textlen + framelen],
	dtype=torch.long,
	device=data.device)
	torch.arange(512 + framelen * 2,
	512 + framelen * 3,
	out=position_ids[textlen + framelen:textlen + framelen * 2],
	dtype=torch.long,
	device=data.device)
	torch.arange(512 + framelen * (frame_num - 1),
	512 + framelen * frame_num,
	out=position_ids[textlen + framelen * 2:textlen +
	framelen * 3],
	dtype=torch.long,
	device=data.device)
	torch.arange(512 + framelen * 1,
	512 + framelen * 2,
	out=position_ids[textlen + framelen * 3:textlen +
	framelen * 4],
	dtype=torch.long,
	device=data.device)
	torch.arange(512 + framelen * 3,
	512 + framelen * 4,
	out=position_ids[textlen + framelen * 4:textlen +
	framelen * 5],
	dtype=torch.long,
	device=data.device)

	position_ids = position_ids.unsqueeze(0)

	return tokens, attention_mask, position_ids


	def my_update_mems(hiddens, mems_buffers, mems_indexs,
	limited_spatial_channel_mem, text_len, frame_len):
	if hiddens is None:
	return None, mems_indexs
	mem_num = len(hiddens)
	ret_mem = []
	with torch.no_grad():
	for id in range(mem_num):
	if hiddens[id][0] is None:
	ret_mem.append(None)
	else:
	if id == 0 and limited_spatial_channel_mem and mems_indexs[
	id] + hiddens[0][0].shape[1] >= text_len + frame_len:
	if mems_indexs[id] == 0:
	for layer, hidden in enumerate(hiddens[id]):
	mems_buffers[id][
	layer, :, :text_len] = hidden.expand(
	mems_buffers[id].shape[1], -1,
	-1)[:, :text_len]
	new_mem_len_part2 = (mems_indexs[id] +
	hiddens[0][0].shape[1] -
	text_len) % frame_len
	if new_mem_len_part2 > 0:
	for layer, hidden in enumerate(hiddens[id]):
	mems_buffers[id][
	layer, :, text_len:text_len +
	new_mem_len_part2] = hidden.expand(
	mems_buffers[id].shape[1], -1,
	-1)[:, -new_mem_len_part2:]
	mems_indexs[id] = text_len + new_mem_len_part2
	else:
	for layer, hidden in enumerate(hiddens[id]):
	mems_buffers[id][layer, :,
	mems_indexs[id]:mems_indexs[id] +
	hidden.shape[1]] = hidden.expand(
	mems_buffers[id].shape[1], -1, -1)
	mems_indexs[id] += hidden.shape[1]
	ret_mem.append(mems_buffers[id][:, :, :mems_indexs[id]])
	return ret_mem, mems_indexs


	def calc_next_tokens_frame_begin_id(text_len, frame_len, total_len):
	# The fisrt token's position id of the frame that the next token belongs to;
	if total_len < text_len:
	return None
	return (total_len - text_len) // frame_len * frame_len + text_len


	def my_filling_sequence(
	model,
	tokenizer,
	args,
	seq,
	batch_size,
	get_masks_and_position_ids,
	text_len,
	frame_len,
	strategy=BaseStrategy(),
	strategy2=BaseStrategy(),
	mems=None,
	log_text_attention_weights=0, # default to 0: no artificial change
	mode_stage1=True,
	enforce_no_swin=False,
	guider_seq=None,
	guider_text_len=0,
	guidance_alpha=1,
	limited_spatial_channel_mem=False, # 空间通道的存储限制在本帧内
	**kw_args):
	'''
	seq: [2, 3, 5, ..., -1(to be generated), -1, ...]
	mems: [num_layers, batch_size, len_mems(index), mem_hidden_size]
	cache, should be first mems.shape[1] parts of context_tokens.
	mems are the first-level citizens here, but we don't assume what is memorized.
	input mems are used when multi-phase generation.
	'''
	if guider_seq is not None:
	logger.debug('Using Guidance In Inference')
	if limited_spatial_channel_mem:
	logger.debug("Limit spatial-channel's mem to current frame")
	assert len(seq.shape) == 2

	# building the initial tokens, attention_mask, and position_ids
	actual_context_length = 0

	while seq[-1][
	actual_context_length] >= 0: # the last seq has least given tokens
	actual_context_length += 1 # [0, context_length-1] are given
	assert actual_context_length > 0
	current_frame_num = (actual_context_length - text_len) // frame_len
	assert current_frame_num >= 0
	context_length = text_len + current_frame_num * frame_len

	tokens, attention_mask, position_ids = get_masks_and_position_ids(
	seq, text_len, frame_len)
	tokens = tokens[..., :context_length]
	input_tokens = tokens.clone()

	if guider_seq is not None:
	guider_index_delta = text_len - guider_text_len
	guider_tokens, guider_attention_mask, guider_position_ids = get_masks_and_position_ids(
	guider_seq, guider_text_len, frame_len)
	guider_tokens = guider_tokens[..., :context_length -
	guider_index_delta]
	guider_input_tokens = guider_tokens.clone()

	for fid in range(current_frame_num):
	input_tokens[:, text_len + 400 * fid] = tokenizer['<start_of_image>']
	if guider_seq is not None:
	guider_input_tokens[:, guider_text_len +
	400 * fid] = tokenizer['<start_of_image>']

	attention_mask = attention_mask.type_as(next(
	model.parameters())) # if fp16
	# initialize generation
	counter = context_length - 1 # Last fixed index is ``counter''
	index = 0 # Next forward starting index, also the length of cache.
	mems_buffers_on_GPU = False
	mems_indexs = [0, 0]
	mems_len = [(400 + 74) if limited_spatial_channel_mem else 5 * 400 + 74,
	5 * 400 + 74]
	mems_buffers = [
	torch.zeros(args.num_layers,
	batch_size,
	mem_len,
	args.hidden_size * 2,
	dtype=next(model.parameters()).dtype)
	for mem_len in mems_len
	]

	if guider_seq is not None:
	guider_attention_mask = guider_attention_mask.type_as(
	next(model.parameters())) # if fp16
	guider_mems_buffers = [
	torch.zeros(args.num_layers,
	batch_size,
	mem_len,
	args.hidden_size * 2,
	dtype=next(model.parameters()).dtype)
	for mem_len in mems_len
	]
	guider_mems_indexs = [0, 0]
	guider_mems = None

	torch.cuda.empty_cache()
	# step-by-step generation
	while counter < len(seq[0]) - 1:
	# we have generated counter+1 tokens
	# Now, we want to generate seq[counter + 1],
	# token[:, index: counter+1] needs forwarding.
	if index == 0:
	group_size = 2 if (input_tokens.shape[0] == batch_size
	and not mode_stage1) else batch_size

	logits_all = None
	for batch_idx in range(0, input_tokens.shape[0], group_size):
	logits, *output_per_layers = model(
	input_tokens[batch_idx:batch_idx + group_size, index:],
	position_ids[..., index:counter + 1],
	attention_mask, # TODO memlen
	mems=mems,
	text_len=text_len,
	frame_len=frame_len,
	counter=counter,
	log_text_attention_weights=log_text_attention_weights,
	enforce_no_swin=enforce_no_swin,
	**kw_args)
	logits_all = torch.cat(
	(logits_all,
	logits), dim=0) if logits_all is not None else logits
	mem_kv01 = [[o['mem_kv'][0] for o in output_per_layers],
	[o['mem_kv'][1] for o in output_per_layers]]
	next_tokens_frame_begin_id = calc_next_tokens_frame_begin_id(
	text_len, frame_len, mem_kv01[0][0].shape[1])
	for id, mem_kv in enumerate(mem_kv01):
	for layer, mem_kv_perlayer in enumerate(mem_kv):
	if limited_spatial_channel_mem and id == 0:
	mems_buffers[id][
	layer, batch_idx:batch_idx + group_size, :
	text_len] = mem_kv_perlayer.expand(
	min(group_size,
	input_tokens.shape[0] - batch_idx), -1,
	-1)[:, :text_len]
	mems_buffers[id][layer, batch_idx:batch_idx+group_size, text_len:text_len+mem_kv_perlayer.shape[1]-next_tokens_frame_begin_id] =\
	mem_kv_perlayer.expand(min(group_size, input_tokens.shape[0]-batch_idx), -1, -1)[:, next_tokens_frame_begin_id:]
	else:
	mems_buffers[id][
	layer, batch_idx:batch_idx +
	group_size, :mem_kv_perlayer.
	shape[1]] = mem_kv_perlayer.expand(
	min(group_size,
	input_tokens.shape[0] - batch_idx), -1,
	-1)
	mems_indexs[0], mems_indexs[1] = mem_kv01[0][0].shape[
	1], mem_kv01[1][0].shape[1]
	if limited_spatial_channel_mem:
	mems_indexs[0] -= (next_tokens_frame_begin_id - text_len)

	mems = [
	mems_buffers[id][:, :, :mems_indexs[id]] for id in range(2)
	]
	logits = logits_all

	# Guider
	if guider_seq is not None:
	guider_logits_all = None
	for batch_idx in range(0, guider_input_tokens.shape[0],
	group_size):
	guider_logits, *guider_output_per_layers = model(
	guider_input_tokens[batch_idx:batch_idx + group_size,
	max(index -
	guider_index_delta, 0):],
	guider_position_ids[
	...,
	max(index - guider_index_delta, 0):counter + 1 -
	guider_index_delta],
	guider_attention_mask,
	mems=guider_mems,
	text_len=guider_text_len,
	frame_len=frame_len,
	counter=counter - guider_index_delta,
	log_text_attention_weights=log_text_attention_weights,
	enforce_no_swin=enforce_no_swin,
	**kw_args)
	guider_logits_all = torch.cat(
	(guider_logits_all, guider_logits), dim=0
	) if guider_logits_all is not None else guider_logits
	guider_mem_kv01 = [[
	o['mem_kv'][0] for o in guider_output_per_layers
	], [o['mem_kv'][1] for o in guider_output_per_layers]]
	for id, guider_mem_kv in enumerate(guider_mem_kv01):
	for layer, guider_mem_kv_perlayer in enumerate(
	guider_mem_kv):
	if limited_spatial_channel_mem and id == 0:
	guider_mems_buffers[id][
	layer, batch_idx:batch_idx + group_size, :
	guider_text_len] = guider_mem_kv_perlayer.expand(
	min(group_size,
	input_tokens.shape[0] - batch_idx),
	-1, -1)[:, :guider_text_len]
	guider_next_tokens_frame_begin_id = calc_next_tokens_frame_begin_id(
	guider_text_len, frame_len,
	guider_mem_kv_perlayer.shape[1])
	guider_mems_buffers[id][layer, batch_idx:batch_idx+group_size, guider_text_len:guider_text_len+guider_mem_kv_perlayer.shape[1]-guider_next_tokens_frame_begin_id] =\
	guider_mem_kv_perlayer.expand(min(group_size, input_tokens.shape[0]-batch_idx), -1, -1)[:, guider_next_tokens_frame_begin_id:]
	else:
	guider_mems_buffers[id][
	layer, batch_idx:batch_idx +
	group_size, :guider_mem_kv_perlayer.
	shape[1]] = guider_mem_kv_perlayer.expand(
	min(group_size,
	input_tokens.shape[0] - batch_idx),
	-1, -1)
	guider_mems_indexs[0], guider_mems_indexs[
	1] = guider_mem_kv01[0][0].shape[1], guider_mem_kv01[
	1][0].shape[1]
	if limited_spatial_channel_mem:
	guider_mems_indexs[0] -= (
	guider_next_tokens_frame_begin_id -
	guider_text_len)
	guider_mems = [
	guider_mems_buffers[id][:, :, :guider_mems_indexs[id]]
	for id in range(2)
	]
	guider_logits = guider_logits_all
	else:
	if not mems_buffers_on_GPU:
	if not mode_stage1:
	torch.cuda.empty_cache()
	for idx, mem in enumerate(mems):
	mems[idx] = mem.to(next(model.parameters()).device)
	if guider_seq is not None:
	for idx, mem in enumerate(guider_mems):
	guider_mems[idx] = mem.to(
	next(model.parameters()).device)
	else:
	torch.cuda.empty_cache()
	for idx, mem_buffer in enumerate(mems_buffers):
	mems_buffers[idx] = mem_buffer.to(
	next(model.parameters()).device)
	mems = [
	mems_buffers[id][:, :, :mems_indexs[id]]
	for id in range(2)
	]
	if guider_seq is not None:
	for idx, guider_mem_buffer in enumerate(
	guider_mems_buffers):
	guider_mems_buffers[idx] = guider_mem_buffer.to(
	next(model.parameters()).device)
	guider_mems = [
	guider_mems_buffers[id]
	[:, :, :guider_mems_indexs[id]] for id in range(2)
	]
	mems_buffers_on_GPU = True

	logits, *output_per_layers = model(
	input_tokens[:, index:],
	position_ids[..., index:counter + 1],
	attention_mask, # TODO memlen
	mems=mems,
	text_len=text_len,
	frame_len=frame_len,
	counter=counter,
	log_text_attention_weights=log_text_attention_weights,
	enforce_no_swin=enforce_no_swin,
	limited_spatial_channel_mem=limited_spatial_channel_mem,
	**kw_args)
	mem_kv0, mem_kv1 = [o['mem_kv'][0] for o in output_per_layers
	], [o['mem_kv'][1] for o in output_per_layers]

	if guider_seq is not None:
	guider_logits, *guider_output_per_layers = model(
	guider_input_tokens[:,
	max(index - guider_index_delta, 0):],
	guider_position_ids[...,
	max(index -
	guider_index_delta, 0):counter +
	1 - guider_index_delta],
	guider_attention_mask,
	mems=guider_mems,
	text_len=guider_text_len,
	frame_len=frame_len,
	counter=counter - guider_index_delta,
	log_text_attention_weights=0,
	enforce_no_swin=enforce_no_swin,
	limited_spatial_channel_mem=limited_spatial_channel_mem,
	**kw_args)
	guider_mem_kv0, guider_mem_kv1 = [
	o['mem_kv'][0] for o in guider_output_per_layers
	], [o['mem_kv'][1] for o in guider_output_per_layers]

	if not mems_buffers_on_GPU:
	torch.cuda.empty_cache()
	for idx, mem_buffer in enumerate(mems_buffers):
	mems_buffers[idx] = mem_buffer.to(
	next(model.parameters()).device)
	if guider_seq is not None:
	for idx, guider_mem_buffer in enumerate(
	guider_mems_buffers):
	guider_mems_buffers[idx] = guider_mem_buffer.to(
	next(model.parameters()).device)
	mems_buffers_on_GPU = True

	mems, mems_indexs = my_update_mems([mem_kv0, mem_kv1],
	mems_buffers, mems_indexs,
	limited_spatial_channel_mem,
	text_len, frame_len)
	if guider_seq is not None:
	guider_mems, guider_mems_indexs = my_update_mems(
	[guider_mem_kv0, guider_mem_kv1], guider_mems_buffers,
	guider_mems_indexs, limited_spatial_channel_mem,
	guider_text_len, frame_len)

	counter += 1
	index = counter

	logits = logits[:, -1].expand(batch_size,
	-1) # [batch size, vocab size]
	tokens = tokens.expand(batch_size, -1)
	if guider_seq is not None:
	guider_logits = guider_logits[:, -1].expand(batch_size, -1)
	guider_tokens = guider_tokens.expand(batch_size, -1)

	if seq[-1][counter].item() < 0:
	# sampling
	guided_logits = guider_logits + (
	logits - guider_logits
	) * guidance_alpha if guider_seq is not None else logits
	if mode_stage1 and counter < text_len + 400:
	tokens, mems = strategy.forward(guided_logits, tokens, mems)
	else:
	tokens, mems = strategy2.forward(guided_logits, tokens, mems)
	if guider_seq is not None:
	guider_tokens = torch.cat((guider_tokens, tokens[:, -1:]),
	dim=1)

	if seq[0][counter].item() >= 0:
	for si in range(seq.shape[0]):
	if seq[si][counter].item() >= 0:
	tokens[si, -1] = seq[si, counter]
	if guider_seq is not None:
	guider_tokens[si,
	-1] = guider_seq[si, counter -
	guider_index_delta]

	else:
	tokens = torch.cat(
	(tokens, seq[:, counter:counter + 1].clone().expand(
	tokens.shape[0], 1).to(device=tokens.device,
	dtype=tokens.dtype)),
	dim=1)
	if guider_seq is not None:
	guider_tokens = torch.cat(
	(guider_tokens,
	guider_seq[:, counter - guider_index_delta:counter + 1 -
	guider_index_delta].clone().expand(
	guider_tokens.shape[0], 1).to(
	device=guider_tokens.device,
	dtype=guider_tokens.dtype)),
	dim=1)

	input_tokens = tokens.clone()
	if guider_seq is not None:
	guider_input_tokens = guider_tokens.clone()
	if (index - text_len - 1) // 400 < (input_tokens.shape[-1] - text_len -
	1) // 400:
	boi_idx = ((index - text_len - 1) // 400 + 1) * 400 + text_len
	while boi_idx < input_tokens.shape[-1]:
	input_tokens[:, boi_idx] = tokenizer['<start_of_image>']
	if guider_seq is not None:
	guider_input_tokens[:, boi_idx -
	guider_index_delta] = tokenizer[
	'<start_of_image>']
	boi_idx += 400

	if strategy.is_done:
	break
	return strategy.finalize(tokens, mems)


	class InferenceModel_Sequential(CogVideoCacheModel):
	def __init__(self, args, transformer=None, parallel_output=True):
	super().__init__(args,
	transformer=transformer,
	parallel_output=parallel_output,
	window_size=-1,
	cogvideo_stage=1)

	# TODO: check it

	def final_forward(self, logits, **kwargs):
	logits_parallel = logits
	logits_parallel = torch.nn.functional.linear(
	logits_parallel.float(),
	self.transformer.word_embeddings.weight[:20000].float())
	return logits_parallel


	class InferenceModel_Interpolate(CogVideoCacheModel):
	def __init__(self, args, transformer=None, parallel_output=True):
	super().__init__(args,
	transformer=transformer,
	parallel_output=parallel_output,
	window_size=10,
	cogvideo_stage=2)

	# TODO: check it

	def final_forward(self, logits, **kwargs):
	logits_parallel = logits
	logits_parallel = torch.nn.functional.linear(
	logits_parallel.float(),
	self.transformer.word_embeddings.weight[:20000].float())
	return logits_parallel


	def get_default_args() -> argparse.Namespace:
	known = argparse.Namespace(generate_frame_num=5,
	coglm_temperature2=0.89,
	use_guidance_stage1=True,
	use_guidance_stage2=False,
	guidance_alpha=3.0,
	stage_1=True,
	stage_2=False,
	both_stages=False,
	parallel_size=1,
	stage1_max_inference_batch_size=-1,
	multi_gpu=False,
	layout='64, 464, 2064',
	window_size=10,
	additional_seqlen=2000,
	cogvideo_stage=1)

	args_list = [
	'--tokenizer-type',
	'fake',
	'--mode',
	'inference',
	'--distributed-backend',
	'nccl',
	'--fp16',
	'--model-parallel-size',
	'1',
	'--temperature',
	'1.05',
	'--top_k',
	'12',
	'--sandwich-ln',
	'--seed',
	'1234',
	'--num-workers',
	'0',
	'--batch-size',
	'1',
	'--max-inference-batch-size',
	'8',
	]
	args = get_args(args_list)
	args = argparse.Namespace(vars(args), vars(known))
	args.layout = [int(x) for x in args.layout.split(',')]
	args.do_train = False
	return args


	class Model:
	def __init__(self, only_first_stage: bool = False):
	self.args = get_default_args()
	if only_first_stage:
	self.args.stage_1 = True
	self.args.both_stages = False
	else:
	self.args.stage_1 = False
	self.args.both_stages = True

	self.tokenizer = self.load_tokenizer()

	self.model_stage1, self.args = self.load_model_stage1()
	self.model_stage2, self.args = self.load_model_stage2()

	self.strategy_cogview2, self.strategy_cogvideo = self.load_strategies()
	self.dsr = self.load_dsr()

	self.device = torch.device(self.args.device)

	def load_tokenizer(self) -> IceTokenizer:
	logger.info('--- load_tokenizer ---')
	start = time.perf_counter()

	tokenizer = IceTokenizer(ICETK_MODEL_DIR.as_posix())
	tokenizer.add_special_tokens(
	['<start_of_image>', '<start_of_english>', '<start_of_chinese>'])

	elapsed = time.perf_counter() - start
	logger.info(f'--- done ({elapsed=:.3f}) ---')
	return tokenizer

	def load_model_stage1(
	self) -> tuple[CogVideoCacheModel, argparse.Namespace]:
	logger.info('--- load_model_stage1 ---')
	start = time.perf_counter()

	args = self.args
	model_stage1, args = InferenceModel_Sequential.from_pretrained(
	args, 'cogvideo-stage1')
	model_stage1.eval()
	if args.both_stages:
	model_stage1 = model_stage1.cpu()

	elapsed = time.perf_counter() - start
	logger.info(f'--- done ({elapsed=:.3f}) ---')
	return model_stage1, args

	def load_model_stage2(
	self) -> tuple[CogVideoCacheModel \| None, argparse.Namespace]:
	logger.info('--- load_model_stage2 ---')
	start = time.perf_counter()

	args = self.args
	if args.both_stages:
	model_stage2, args = InferenceModel_Interpolate.from_pretrained(
	args, 'cogvideo-stage2')
	model_stage2.eval()
	if args.both_stages:
	model_stage2 = model_stage2.cpu()
	else:
	model_stage2 = None

	elapsed = time.perf_counter() - start
	logger.info(f'--- done ({elapsed=:.3f}) ---')
	return model_stage2, args

	def load_strategies(self) -> tuple[CoglmStrategy, CoglmStrategy]:
	logger.info('--- load_strategies ---')
	start = time.perf_counter()

	invalid_slices = [slice(self.tokenizer.num_image_tokens, None)]
	strategy_cogview2 = CoglmStrategy(invalid_slices,
	temperature=1.0,
	top_k=16)
	strategy_cogvideo = CoglmStrategy(
	invalid_slices,
	temperature=self.args.temperature,
	top_k=self.args.top_k,
	temperature2=self.args.coglm_temperature2)

	elapsed = time.perf_counter() - start
	logger.info(f'--- done ({elapsed=:.3f}) ---')
	return strategy_cogview2, strategy_cogvideo

	def load_dsr(self) -> DirectSuperResolution \| None:
	logger.info('--- load_dsr ---')
	start = time.perf_counter()

	if self.args.both_stages:
	path = auto_create('cogview2-dsr', path=None)
	dsr = DirectSuperResolution(self.args,
	path,
	max_bz=12,
	onCUDA=False)
	else:
	dsr = None

	elapsed = time.perf_counter() - start
	logger.info(f'--- done ({elapsed=:.3f}) ---')
	return dsr

	@torch.inference_mode()
	def process_stage1(self,
	model,
	seq_text,
	duration,
	video_raw_text=None,
	video_guidance_text='视频',
	image_text_suffix='',
	batch_size=1):
	process_start_time = time.perf_counter()

	generate_frame_num = self.args.generate_frame_num
	tokenizer = self.tokenizer
	use_guide = self.args.use_guidance_stage1

	if next(model.parameters()).device != self.device:
	move_start_time = time.perf_counter()
	logger.debug('moving stage 1 model to cuda')

	model = model.to(self.device)

	elapsed = time.perf_counter() - move_start_time
	logger.debug(f'moving in model1 takes time: {elapsed:.2f}')

	if video_raw_text is None:
	video_raw_text = seq_text
	mbz = self.args.stage1_max_inference_batch_size if self.args.stage1_max_inference_batch_size > 0 else self.args.max_inference_batch_size
	assert batch_size < mbz or batch_size % mbz == 0
	frame_len = 400

	# generate the first frame:
	enc_text = tokenizer.encode(seq_text + image_text_suffix)
	seq_1st = enc_text + [tokenizer['<start_of_image>']] + [-1] * 400
	logger.info(
	f'[Generating First Frame with CogView2] Raw text: {tokenizer.decode(enc_text):s}'
	)
	text_len_1st = len(seq_1st) - frame_len * 1 - 1

	seq_1st = torch.tensor(seq_1st, dtype=torch.long,
	device=self.device).unsqueeze(0)
	output_list_1st = []
	for tim in range(max(batch_size // mbz, 1)):
	start_time = time.perf_counter()
	output_list_1st.append(
	my_filling_sequence(
	model,
	tokenizer,
	self.args,
	seq_1st.clone(),
	batch_size=min(batch_size, mbz),
	get_masks_and_position_ids=
	get_masks_and_position_ids_stage1,
	text_len=text_len_1st,
	frame_len=frame_len,
	strategy=self.strategy_cogview2,
	strategy2=self.strategy_cogvideo,
	log_text_attention_weights=1.4,
	enforce_no_swin=True,
	mode_stage1=True,
	)[0])
	elapsed = time.perf_counter() - start_time
	logger.info(f'[First Frame] Elapsed: {elapsed:.2f}')
	output_tokens_1st = torch.cat(output_list_1st, dim=0)
	given_tokens = output_tokens_1st[:, text_len_1st + 1:text_len_1st +
	401].unsqueeze(
	1
	) # given_tokens.shape: [bs, frame_num, 400]

	# generate subsequent frames:
	total_frames = generate_frame_num
	enc_duration = tokenizer.encode(f'{float(duration)}秒')
	if use_guide:
	video_raw_text = video_raw_text + ' 视频'
	enc_text_video = tokenizer.encode(video_raw_text)
	seq = enc_duration + [tokenizer['<n>']] + enc_text_video + [
	tokenizer['<start_of_image>']
	] + [-1] * 400 * generate_frame_num
	guider_seq = enc_duration + [tokenizer['<n>']] + tokenizer.encode(
	video_guidance_text) + [tokenizer['<start_of_image>']
	] + [-1] * 400 * generate_frame_num
	logger.info(
	f'[Stage1: Generating Subsequent Frames, Frame Rate {4/duration:.1f}] raw text: {tokenizer.decode(enc_text_video):s}'
	)

	text_len = len(seq) - frame_len * generate_frame_num - 1
	guider_text_len = len(guider_seq) - frame_len * generate_frame_num - 1
	seq = torch.tensor(seq, dtype=torch.long,
	device=self.device).unsqueeze(0).repeat(
	batch_size, 1)
	guider_seq = torch.tensor(guider_seq,
	dtype=torch.long,
	device=self.device).unsqueeze(0).repeat(
	batch_size, 1)

	for given_frame_id in range(given_tokens.shape[1]):
	seq[:, text_len + 1 + given_frame_id * 400:text_len + 1 +
	(given_frame_id + 1) * 400] = given_tokens[:, given_frame_id]
	guider_seq[:, guider_text_len + 1 +
	given_frame_id * 400:guider_text_len + 1 +
	(given_frame_id + 1) *
	400] = given_tokens[:, given_frame_id]
	output_list = []

	if use_guide:
	video_log_text_attention_weights = 0
	else:
	guider_seq = None
	video_log_text_attention_weights = 1.4

	for tim in range(max(batch_size // mbz, 1)):
	input_seq = seq[:min(batch_size, mbz)].clone(
	) if tim == 0 else seq[mbz * tim:mbz * (tim + 1)].clone()
	guider_seq2 = (guider_seq[:min(batch_size, mbz)].clone()
	if tim == 0 else guider_seq[mbz * tim:mbz *
	(tim + 1)].clone()
	) if guider_seq is not None else None
	output_list.append(
	my_filling_sequence(
	model,
	tokenizer,
	self.args,
	input_seq,
	batch_size=min(batch_size, mbz),
	get_masks_and_position_ids=
	get_masks_and_position_ids_stage1,
	text_len=text_len,
	frame_len=frame_len,
	strategy=self.strategy_cogview2,
	strategy2=self.strategy_cogvideo,
	log_text_attention_weights=video_log_text_attention_weights,
	guider_seq=guider_seq2,
	guider_text_len=guider_text_len,
	guidance_alpha=self.args.guidance_alpha,
	limited_spatial_channel_mem=True,
	mode_stage1=True,
	)[0])

	output_tokens = torch.cat(output_list, dim=0)[:, 1 + text_len:]

	if self.args.both_stages:
	move_start_time = time.perf_counter()
	logger.debug('moving stage 1 model to cpu')
	model = model.cpu()
	torch.cuda.empty_cache()
	elapsed = time.perf_counter() - move_start_time
	logger.debug(f'moving in model1 takes time: {elapsed:.2f}')

	# decoding
	res = []
	for seq in output_tokens:
	decoded_imgs = [
	self.postprocess(
	torch.nn.functional.interpolate(tokenizer.decode(
	image_ids=seq.tolist()[i * 400:(i + 1) * 400]),
	size=(480, 480))[0])
	for i in range(total_frames)
	]
	res.append(decoded_imgs) # only the last image (target)

	assert len(res) == batch_size
	tokens = output_tokens[:, :+total_frames * 400].reshape(
	-1, total_frames, 400).cpu()

	elapsed = time.perf_counter() - process_start_time
	logger.info(f'--- done ({elapsed=:.3f}) ---')
	return tokens, res[0]

	@torch.inference_mode()
	def process_stage2(self,
	model,
	seq_text,
	duration,
	parent_given_tokens,
	video_raw_text=None,
	video_guidance_text='视频',
	gpu_rank=0,
	gpu_parallel_size=1):
	process_start_time = time.perf_counter()

	generate_frame_num = self.args.generate_frame_num
	tokenizer = self.tokenizer
	use_guidance = self.args.use_guidance_stage2

	stage2_start_time = time.perf_counter()

	if next(model.parameters()).device != self.device:
	move_start_time = time.perf_counter()
	logger.debug('moving stage-2 model to cuda')

	model = model.to(self.device)

	elapsed = time.perf_counter() - move_start_time
	logger.debug(f'moving in stage-2 model takes time: {elapsed:.2f}')

	try:
	sample_num_allgpu = parent_given_tokens.shape[0]
	sample_num = sample_num_allgpu // gpu_parallel_size
	assert sample_num * gpu_parallel_size == sample_num_allgpu
	parent_given_tokens = parent_given_tokens[gpu_rank *
	sample_num:(gpu_rank +
	1) *
	sample_num]
	except:
	logger.critical('No frame_tokens found in interpolation, skip')
	return False, []

	# CogVideo Stage2 Generation
	while duration >= 0.5: # TODO: You can change the boundary to change the frame rate
	parent_given_tokens_num = parent_given_tokens.shape[1]
	generate_batchsize_persample = (parent_given_tokens_num - 1) // 2
	generate_batchsize_total = generate_batchsize_persample * sample_num
	total_frames = generate_frame_num
	frame_len = 400
	enc_text = tokenizer.encode(seq_text)
	enc_duration = tokenizer.encode(str(float(duration)) + '秒')
	seq = enc_duration + [tokenizer['<n>']] + enc_text + [
	tokenizer['<start_of_image>']
	] + [-1] * 400 * generate_frame_num
	text_len = len(seq) - frame_len * generate_frame_num - 1

	logger.info(
	f'[Stage2: Generating Frames, Frame Rate {int(4/duration):d}] raw text: {tokenizer.decode(enc_text):s}'
	)

	# generation
	seq = torch.tensor(seq, dtype=torch.long,
	device=self.device).unsqueeze(0).repeat(
	generate_batchsize_total, 1)
	for sample_i in range(sample_num):
	for i in range(generate_batchsize_persample):
	seq[sample_i * generate_batchsize_persample +
	i][text_len + 1:text_len + 1 +
	400] = parent_given_tokens[sample_i][2 * i]
	seq[sample_i * generate_batchsize_persample +
	i][text_len + 1 + 400:text_len + 1 +
	800] = parent_given_tokens[sample_i][2 * i + 1]
	seq[sample_i * generate_batchsize_persample +
	i][text_len + 1 + 800:text_len + 1 +
	1200] = parent_given_tokens[sample_i][2 * i + 2]

	if use_guidance:
	guider_seq = enc_duration + [
	tokenizer['<n>']
	] + tokenizer.encode(video_guidance_text) + [
	tokenizer['<start_of_image>']
	] + [-1] * 400 * generate_frame_num
	guider_text_len = len(
	guider_seq) - frame_len * generate_frame_num - 1
	guider_seq = torch.tensor(
	guider_seq, dtype=torch.long,
	device=self.device).unsqueeze(0).repeat(
	generate_batchsize_total, 1)
	for sample_i in range(sample_num):
	for i in range(generate_batchsize_persample):
	guider_seq[sample_i * generate_batchsize_persample +
	i][text_len + 1:text_len + 1 +
	400] = parent_given_tokens[sample_i][2 *
	i]
	guider_seq[sample_i * generate_batchsize_persample +
	i][text_len + 1 + 400:text_len + 1 +
	800] = parent_given_tokens[sample_i][2 *
	i +
	1]
	guider_seq[sample_i * generate_batchsize_persample +
	i][text_len + 1 + 800:text_len + 1 +
	1200] = parent_given_tokens[sample_i][2 *
	i +
	2]
	video_log_text_attention_weights = 0
	else:
	guider_seq = None
	guider_text_len = 0
	video_log_text_attention_weights = 1.4

	mbz = self.args.max_inference_batch_size

	assert generate_batchsize_total < mbz or generate_batchsize_total % mbz == 0
	output_list = []
	start_time = time.perf_counter()
	for tim in range(max(generate_batchsize_total // mbz, 1)):
	input_seq = seq[:min(generate_batchsize_total, mbz)].clone(
	) if tim == 0 else seq[mbz * tim:mbz * (tim + 1)].clone()
	guider_seq2 = (
	guider_seq[:min(generate_batchsize_total, mbz)].clone()
	if tim == 0 else guider_seq[mbz * tim:mbz *
	(tim + 1)].clone()
	) if guider_seq is not None else None
	output_list.append(
	my_filling_sequence(
	model,
	tokenizer,
	self.args,
	input_seq,
	batch_size=min(generate_batchsize_total, mbz),
	get_masks_and_position_ids=
	get_masks_and_position_ids_stage2,
	text_len=text_len,
	frame_len=frame_len,
	strategy=self.strategy_cogview2,
	strategy2=self.strategy_cogvideo,
	log_text_attention_weights=
	video_log_text_attention_weights,
	mode_stage1=False,
	guider_seq=guider_seq2,
	guider_text_len=guider_text_len,
	guidance_alpha=self.args.guidance_alpha,
	limited_spatial_channel_mem=True,
	)[0])
	elapsed = time.perf_counter() - start_time
	logger.info(f'Duration {duration:.2f}, Elapsed: {elapsed:.2f}\n')

	output_tokens = torch.cat(output_list, dim=0)
	output_tokens = output_tokens[:, text_len + 1:text_len + 1 +
	(total_frames) * 400].reshape(
	sample_num, -1,
	400 * total_frames)
	output_tokens_merge = torch.cat(
	(output_tokens[:, :, :1 * 400], output_tokens[:, :,
	400 * 3:4 * 400],
	output_tokens[:, :, 400 * 1:2 * 400],
	output_tokens[:, :, 400 * 4:(total_frames) * 400]),
	dim=2).reshape(sample_num, -1, 400)

	output_tokens_merge = torch.cat(
	(output_tokens_merge, output_tokens[:, -1:, 400 * 2:3 * 400]),
	dim=1)
	duration /= 2
	parent_given_tokens = output_tokens_merge

	if self.args.both_stages:
	move_start_time = time.perf_counter()
	logger.debug('moving stage 2 model to cpu')
	model = model.cpu()
	torch.cuda.empty_cache()
	elapsed = time.perf_counter() - move_start_time
	logger.debug(f'moving out model2 takes time: {elapsed:.2f}')

	elapsed = time.perf_counter() - stage2_start_time
	logger.info(f'CogVideo Stage2 completed. Elapsed: {elapsed:.2f}\n')

	# direct super-resolution by CogView2
	logger.info('[Direct super-resolution]')
	dsr_start_time = time.perf_counter()

	enc_text = tokenizer.encode(seq_text)
	frame_num_per_sample = parent_given_tokens.shape[1]
	parent_given_tokens_2d = parent_given_tokens.reshape(-1, 400)
	text_seq = torch.tensor(enc_text, dtype=torch.long,
	device=self.device).unsqueeze(0).repeat(
	parent_given_tokens_2d.shape[0], 1)
	sred_tokens = self.dsr(text_seq, parent_given_tokens_2d)

	decoded_sr_videos = []
	for sample_i in range(sample_num):
	decoded_sr_imgs = []
	for frame_i in range(frame_num_per_sample):
	decoded_sr_img = tokenizer.decode(
	image_ids=sred_tokens[frame_i + sample_i *
	frame_num_per_sample][-3600:])
	decoded_sr_imgs.append(
	self.postprocess(
	torch.nn.functional.interpolate(decoded_sr_img,
	size=(480, 480))[0]))
	decoded_sr_videos.append(decoded_sr_imgs)

	elapsed = time.perf_counter() - dsr_start_time
	logger.info(
	f'Direct super-resolution completed. Elapsed: {elapsed:.2f}')

	elapsed = time.perf_counter() - process_start_time
	logger.info(f'--- done ({elapsed=:.3f}) ---')
	return True, decoded_sr_videos[0]

	@staticmethod
	def postprocess(tensor: torch.Tensor) -> np.ndarray:
	return tensor.cpu().mul(255).add_(0.5).clamp_(0, 255).permute(
	1, 2, 0).to(torch.uint8).numpy()

	def run(self, text: str, seed: int,
	only_first_stage: bool) -> list[np.ndarray]:
	logger.info('==================== run ====================')
	start = time.perf_counter()

	set_random_seed(seed)

	if only_first_stage:
	self.args.stage_1 = True
	self.args.both_stages = False
	else:
	self.args.stage_1 = False
	self.args.both_stages = True

	parent_given_tokens, res = self.process_stage1(
	self.model_stage1,
	text,
	duration=4.0,
	video_raw_text=text,
	video_guidance_text='视频',
	image_text_suffix=' 高清摄影',
	batch_size=self.args.batch_size)
	if not only_first_stage:
	_, res = self.process_stage2(
	self.model_stage2,
	text,
	duration=2.0,
	parent_given_tokens=parent_given_tokens,
	video_raw_text=text + ' 视频',
	video_guidance_text='视频',
	gpu_rank=0,
	gpu_parallel_size=1) # TODO: 修改

	elapsed = time.perf_counter() - start
	logger.info(f'Elapsed: {elapsed:.3f}')
	logger.info('==================== done ====================')
	return res


	class AppModel(Model):
	def __init__(self, only_first_stage: bool):
	super().__init__(only_first_stage)
	self.translator = gr.Interface.load(
	'spaces/chinhon/translation_eng2ch')

	def to_video(self, frames: list[np.ndarray]) -> str:
	out_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
	if self.args.stage_1:
	fps = 4
	else:
	fps = 8
	writer = iio.get_writer(out_file.name, fps=fps)
	for frame in frames:
	writer.append_data(frame)
	writer.close()
	return out_file.name

	def run_with_translation(
	self, text: str, translate: bool, seed: int, only_first_stage: bool
	) -> tuple[str \| None, np.ndarray \| None, list[np.ndarray] \| None]:
	logger.info(f'{text=}, {translate=}, {seed=}, {only_first_stage=}')
	if translate:
	text = translated_text = self.translator(text)
	else:
	translated_text = None
	frames = self.run(text, seed, only_first_stage)
	video_path = self.to_video(frames)
	return translated_text, video_path, frames