(self, args, **kwargs)
| 29 | |
| 30 | class SATVideoDiffusionEngine(nn.Module): |
| 31 | def __init__(self, args, **kwargs): |
| 32 | super().__init__() |
| 33 | |
| 34 | model_config = args.model_config |
| 35 | # model args preprocess |
| 36 | log_keys = model_config.get("log_keys", None) |
| 37 | input_key = model_config.get("input_key", "mp4") |
| 38 | network_config = model_config.get("network_config", None) |
| 39 | network_wrapper = model_config.get("network_wrapper", None) |
| 40 | denoiser_config = model_config.get("denoiser_config", None) |
| 41 | sampler_config = model_config.get("sampler_config", None) |
| 42 | conditioner_config = model_config.get("conditioner_config", None) |
| 43 | first_stage_config = model_config.get("first_stage_config", None) |
| 44 | loss_fn_config = model_config.get("loss_fn_config", None) |
| 45 | scale_factor = model_config.get("scale_factor", 1.0) |
| 46 | latent_input = model_config.get("latent_input", False) |
| 47 | disable_first_stage_autocast = model_config.get("disable_first_stage_autocast", False) |
| 48 | no_cond_log = model_config.get("disable_first_stage_autocast", False) |
| 49 | not_trainable_prefixes = model_config.get("not_trainable_prefixes", ["first_stage_model", "conditioner"]) |
| 50 | compile_model = model_config.get("compile_model", False) |
| 51 | en_and_decode_n_samples_a_time = model_config.get("en_and_decode_n_samples_a_time", None) |
| 52 | en_and_decode_n_frames_a_time = model_config.get("en_and_decode_n_frames_a_time", None) |
| 53 | lr_scale = model_config.get("lr_scale", None) |
| 54 | lora_train = model_config.get("lora_train", False) |
| 55 | self.use_pd = model_config.get("use_pd", False) # progressive distillation |
| 56 | |
| 57 | self.cond_inds = model_config.get("cond_inds_sampling", None) |
| 58 | |
| 59 | self.log_keys = log_keys |
| 60 | self.input_key = input_key |
| 61 | self.not_trainable_prefixes = not_trainable_prefixes |
| 62 | self.en_and_decode_n_samples_a_time = en_and_decode_n_samples_a_time |
| 63 | self.en_and_decode_n_frames_a_time = en_and_decode_n_frames_a_time |
| 64 | |
| 65 | self.truncate_n_frames_decode = model_config.get("truncate_n_frames_decode", None) |
| 66 | # * 13 latents for 49 frames |
| 67 | # * 8 latents for 33 frames |
| 68 | # * 6 latents for 25 frames |
| 69 | |
| 70 | self.lr_scale = lr_scale |
| 71 | self.lora_train = lora_train |
| 72 | self.noised_image_input = model_config.get("noised_image_input", False) |
| 73 | self.noised_image_all_concat = model_config.get("noised_image_all_concat", False) |
| 74 | self.noised_image_dropout = model_config.get("noised_image_dropout", 0.0) |
| 75 | if args.fp16: |
| 76 | dtype = torch.float16 |
| 77 | dtype_str = "fp16" |
| 78 | elif args.bf16: |
| 79 | dtype = torch.bfloat16 |
| 80 | dtype_str = "bf16" |
| 81 | else: |
| 82 | dtype = torch.float32 |
| 83 | dtype_str = "fp32" |
| 84 | self.dtype = dtype |
| 85 | self.dtype_str = dtype_str |
| 86 | |
| 87 | network_config["params"]["dtype"] = dtype_str |
| 88 | model = instantiate_from_config(network_config) |
nothing calls this directly
no test coverage detected