JLNTV_IPD/CogVideo
4
0
Fork 0
mirror of https://github.com/THUDM/CogVideo.git synced 2026-06-01 09:04:08 +08:00
Code Issues Packages Projects Releases Wiki Activity

revert refactor

Browse Source
This commit is contained in:
Stephan Auerhahn 2022-07-25 11:12:18 +00:00
parent 4d94efa6d4
commit 8cc364086b
1 changed files with 83 additions and 73 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

156
predict.py
View File

@ -485,34 +485,25 @@ class Predictor(BasePredictor):
) -> typing.Iterator[Path]:
if translate:
prompt = self.translator.translate(prompt.strip())
if seed == -1:
seed = randint(0, 2**32)
self.args.seed = seed
self.args.use_guidance_stage1 = use_guidance
self.prompt = prompt
self.args.both_stages = both_stages
yield [Path(file) for file in self.run()]
logging.debug("complete, exiting")
return
def run(self):
invalid_slices = [slice(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)
torch.manual_seed(self.args.seed)
random.seed(self.args.seed)
workdir = tempfile.mkdtemp()
os.makedirs(f"{workdir}/output/stage1", exist_ok=True)
os.makedirs(f"{workdir}/output/stage2", exist_ok=True)
if seed == -1:
seed = randint(0, 2**32)
invalid_slices = [slice(tokenizer.num_image_tokens, None)]
self.strategy_cogview2 = CoglmStrategy(invalid_slices,
temperature=1.0, top_k=16)
self.strategy_cogvideo = CoglmStrategy(invalid_slices,
temperature=self.args.temperature, top_k=self.args.top_k,
temperature2=self.args.coglm_temperature2)
torch.manual_seed(seed)
random.seed(seed)
self.args.seed = seed
self.args.use_guidance_stage1 = use_guidance
move_start_time = time.time()
logging.debug("moving stage 2 model to cpu")
self.model_stage2 = self.model_stage2.cpu()
@ -520,17 +511,35 @@ class Predictor(BasePredictor):
logging.debug("moving stage 1 model to cuda")
self.model_stage1 = self.model_stage1.cuda()
logging.debug("moving in model1 takes time: {:.2f}".format(time.time()-move_start_time))
parent_given_tokens = self.process_stage1(self.model_stage1, prompt, duration=4.0, video_raw_text=prompt, video_guidance_text="视频",
image_text_suffix=" 高清摄影",
outputdir=f'{workdir}/output/stage1', batch_size=1)
logging.debug("moving stage 1 model to cpu")
self.model_stage1 = self.model_stage1.cpu()
torch.cuda.empty_cache()
yield Path(f"{workdir}/output/stage1/0.gif")
if both_stages:
move_start_time = time.time()
logging.debug("moving stage 2 model to cuda")
self.model_stage2 = self.model_stage2.cuda()
logging.debug("moving in model2 takes time: {:.2f}".format(time.time()-move_start_time))
self.process_stage2(self.model_stage2, prompt, duration=2.0, video_raw_text=prompt+" 视频",
video_guidance_text="视频", parent_given_tokens=parent_given_tokens,
outputdir=f'{workdir}/output/stage2',
gpu_rank=0, gpu_parallel_size=1)
yield Path(f"{workdir}/output/stage2/0.gif")
logging.debug("complete, exiting")
def process_stage1(self, model, seq_text, duration, video_raw_text=None, video_guidance_text="视频", image_text_suffix="", outputdir=None, batch_size=1):
process_start_time = time.time()
args = self.args
use_guide = args.use_guidance_stage1
batch_size = 1
seq_text = self.prompt
video_raw_text = self.prompt
duration=4.0
video_guidance_text="视频"
image_text_suffix=" 高清摄影"
outputdir=f'{workdir}/output/stage1'
if video_raw_text is None:
video_raw_text = seq_text
mbz = args.stage1_max_inference_batch_size if args.stage1_max_inference_batch_size > 0 else args.max_inference_batch_size
assert batch_size < mbz or batch_size % mbz == 0
frame_len = 400
@ -546,13 +555,13 @@ class Predictor(BasePredictor):
for tim in range(max(batch_size // mbz, 1)):
start_time = time.time()
output_list_1st.append(
my_filling_sequence(self.model_stage1, args,seq_1st.clone(),
my_filling_sequence(model, 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=strategy_cogview2,
strategy2=strategy_cogvideo,
strategy=self.strategy_cogview2,
strategy2=self.strategy_cogvideo,
log_text_attention_weights=1.4,
enforce_no_swin=True,
mode_stage1=True,
@ -593,12 +602,12 @@ class Predictor(BasePredictor):
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(self.model_stage1, args,input_seq,
my_filling_sequence(model, 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=strategy_cogview2,
strategy2=strategy_cogvideo,
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,
@ -622,37 +631,27 @@ class Predictor(BasePredictor):
for clip_i in range(len(imgs)):
# os.makedirs(output_dir_full_paths[clip_i], exist_ok=True)
my_save_multiple_images(imgs[clip_i], outputdir, subdir=f"frames/{clip_i}", debug=False)
out_filename = f'{outputdir}/{clip_i}.gif'
os.system(f"gifmaker -i '{outputdir}'/frames/'{clip_i}'/0*.jpg -o '{out_filename}' -d 0.25")
yield out_filename
os.system(f"gifmaker -i '{outputdir}'/frames/'{clip_i}'/0*.jpg -o '{outputdir}/{clip_i}.gif' -d 0.25")
torch.save(save_tokens, os.path.join(outputdir, 'frame_tokens.pt'))
logging.info("CogVideo Stage1 completed. Taken time {:.2f}\n".format(time.time() - process_start_time))
logging.debug("moving stage 1 model to cpu")
self.model_stage1 = self.model_stage1.cpu()
torch.cuda.empty_cache()
return save_tokens
if not self.args.both_stages:
logging.info("only stage 1 selected, exiting")
return
gpu_rank=0
gpu_parallel_size=1
video_raw_text=self.prompt+" 视频"
duration=2.0
video_guidance_text="视频"
outputdir=f'{workdir}/output/stage2'
parent_given_tokens = save_tokens
stage2_starttime = time.time()
use_guidance = args.use_guidance_stage2
move_start_time = time.time()
logging.debug("moving stage-2 model to cuda")
self.model_stage2 = self.model_stage2.cuda()
logging.debug("moving in stage-2 model takes time: {:.2f}".format(time.time()-move_start_time))
def process_stage2(self, model, seq_text, duration, video_raw_text=None, video_guidance_text="视频", parent_given_tokens=None, conddir=None, outputdir=None, gpu_rank=0, gpu_parallel_size=1):
stage2_starttime = time.time()
args = self.args
use_guidance = args.use_guidance_stage2
if args.both_stages:
move_start_time = time.time()
logging.debug("moving stage-2 model to cuda")
model = model.cuda()
logging.debug("moving in stage-2 model takes time: {:.2f}".format(time.time()-move_start_time))
try:
if parent_given_tokens is None:
assert conddir is not None
parent_given_tokens = torch.load(os.path.join(conddir, 'frame_tokens.pt'), map_location='cpu')
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
@ -707,12 +706,12 @@ class Predictor(BasePredictor):
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(self.model_stage2, args, input_seq,
my_filling_sequence(model, 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=strategy_cogview2,
strategy2=strategy_cogvideo,
strategy=self.strategy_cogview2,
strategy2=self.strategy_cogvideo,
log_text_attention_weights=video_log_text_attention_weights,
mode_stage1=False,
guider_seq=guider_seq2,
@ -741,7 +740,8 @@ class Predictor(BasePredictor):
parent_given_tokens_2d = parent_given_tokens.reshape(-1, 400)
logging.debug("moving stage 2 model to cpu")
self.model_stage2 = self.model_stage2.cpu()
self.model_stage2 = self.model_stage2.cpu()
model = model.cpu()
torch.cuda.empty_cache()
# use dsr if loaded
@ -762,9 +762,7 @@ class Predictor(BasePredictor):
for sample_i in range(sample_num):
my_save_multiple_images(decoded_sr_videos[sample_i], outputdir,subdir=f"frames/{sample_i+sample_num*gpu_rank}", debug=False)
output_file = f'{outputdir}/{sample_i+sample_num*gpu_rank}.gif'
os.system(f"gifmaker -i '{outputdir}'/frames/'{sample_i+sample_num*gpu_rank}'/0*.jpg -o '{output_file}' -d 0.125")
yield output_file
os.system(f"gifmaker -i '{outputdir}'/frames/'{sample_i+sample_num*gpu_rank}'/0*.jpg -o '{outputdir}/{sample_i+sample_num*gpu_rank}.gif' -d 0.125")
logging.info("Direct super-resolution completed. Taken time {:.2f}\n".format(time.time() - dsr_starttime))
else:
@ -786,6 +784,18 @@ class Predictor(BasePredictor):
for sample_i in range(sample_num):
my_save_multiple_images(decoded_videos[sample_i], outputdir,subdir=f"frames/{sample_i+sample_num*gpu_rank}", debug=False)
output_file = f'{outputdir}/{sample_i+sample_num*gpu_rank}.gif'
os.system(f"gifmaker -i '{outputdir}'/frames/'{sample_i+sample_num*gpu_rank}'/0*.jpg -o '{output_file}' -d 0.125")
yield output_file
os.system(f"gifmaker -i '{outputdir}'/frames/'{sample_i+sample_num*gpu_rank}'/0*.jpg -o '{outputdir}/{sample_i+sample_num*gpu_rank}.gif' -d 0.125")
#imgs = []
#for seq in output_tokens:
# decoded_imgs = [torch.nn.functional.interpolate(tokenizer.decode(image_ids=seq.tolist()[i*400: (i+1)*400]), size=(480, 480)) for i in range(total_frames)]
# decoded_video.extend(decoded_imgs) # only the last image (target)
#for sample_i in range(sample_num):
# my_save_multiple_images(decoded_video, outputdir,subdir=f"frames/{sample_i+sample_num*gpu_rank}", debug=False)
# os.system(f"gifmaker -i '{outputdir}'/frames/'{sample_i+sample_num*gpu_rank}'/0*.jpg -o '{outputdir}/{sample_i+sample_num*gpu_rank}.gif' -d 0.125")
return True