harry900000 commited on
Commit
0bcab0b
·
1 Parent(s): dcc4583

add try-catch to make sure the log file can be downloaded

Browse files
Files changed (1) hide show
  1. helper.py +42 -37
helper.py CHANGED
@@ -392,48 +392,53 @@ def generate_video_fun(checkpoints_path: str):
392
 
393
  log.info(f"actual_seed: {actual_seed}")
394
 
395
- if rgb_video_path is None or not os.path.isfile(rgb_video_path):
396
- log.warning(f"File `{rgb_video_path}` does not exist")
397
- rgb_video_path = ""
398
-
399
- # add timer to calculate the generation time
400
- start_time = time.time()
401
-
402
- # parse generation configs
403
- args, control_inputs = parse_arguments(
404
- controlnet_specs_in={
405
- "hdmap": {"control_weight": 0.3, "input_control": hdmap_video_input},
406
- "lidar": {"control_weight": 0.7, "input_control": lidar_video_input},
407
- },
408
- input_video_path=rgb_video_path,
409
- checkpoint_dir=checkpoints_path,
410
- prompt=prompt,
411
- negative_prompt=negative_prompt,
412
- sigma_max=80,
413
- offload_text_encoder_model=True,
414
- is_av_sample=True,
415
- num_gpus=1,
416
- seed=seed,
417
- )
 
418
 
419
- # watch gpu memory
420
- watcher = watch_gpu_memory(10, lambda x: log.debug(f"GPU memory (used, total): {x} (MiB)"))
421
 
422
- # start inference
423
- if chunking <= 0:
424
- chunking = None
425
- videos, prompts = inference(args, control_inputs, chunking)
426
 
427
- # print the generation time
428
- end_time = time.time()
429
- log.info(f"Time taken: {end_time - start_time} s")
430
 
431
- # stop the watcher
432
- stop_watcher()
433
 
434
- video = videos[0]
435
 
436
- log.logger.remove(log_handler)
437
- return video, create_zip_for_download(filename=logfile_path, files_to_zip=[video, logfile_path]), actual_seed
 
 
 
 
438
 
439
  return generate_video
 
392
 
393
  log.info(f"actual_seed: {actual_seed}")
394
 
395
+ try:
396
+ if rgb_video_path is None or not os.path.isfile(rgb_video_path):
397
+ log.warning(f"File `{rgb_video_path}` does not exist")
398
+ rgb_video_path = ""
399
+
400
+ # add timer to calculate the generation time
401
+ start_time = time.time()
402
+
403
+ # parse generation configs
404
+ args, control_inputs = parse_arguments(
405
+ controlnet_specs_in={
406
+ "hdmap": {"control_weight": 0.3, "input_control": hdmap_video_input},
407
+ "lidar": {"control_weight": 0.7, "input_control": lidar_video_input},
408
+ },
409
+ input_video_path=rgb_video_path,
410
+ checkpoint_dir=checkpoints_path,
411
+ prompt=prompt,
412
+ negative_prompt=negative_prompt,
413
+ sigma_max=80,
414
+ offload_text_encoder_model=True,
415
+ is_av_sample=True,
416
+ num_gpus=1,
417
+ seed=seed,
418
+ )
419
 
420
+ # watch gpu memory
421
+ watcher = watch_gpu_memory(10, lambda x: log.debug(f"GPU memory (used, total): {x} (MiB)"))
422
 
423
+ # start inference
424
+ if chunking <= 0:
425
+ chunking = None
426
+ videos, prompts = inference(args, control_inputs, chunking)
427
 
428
+ # print the generation time
429
+ end_time = time.time()
430
+ log.info(f"Time taken: {end_time - start_time} s")
431
 
432
+ # stop the watcher
433
+ stop_watcher()
434
 
435
+ video = videos[0]
436
 
437
+ log.logger.remove(log_handler)
438
+ return video, create_zip_for_download(filename=logfile_path, files_to_zip=[video, logfile_path]), actual_seed
439
+ except Exception as e:
440
+ log.logger.remove(log_handler)
441
+ log.exception(e)
442
+ return "", logfile_path, actual_seed
443
 
444
  return generate_video