Method __init__

diffsynth/models/stepvideo_vae.py:909–946

(self,
        in_channels=3,
        out_channels=3,
        z_channels=64,
        num_res_blocks=2,
        model_path=None,
        weight_dict={},
        world_size=1,
        version=2,
    )

class StepVideoVAE(nn.Module):
    def __init__(self,
        in_channels=3,
        out_channels=3,
        z_channels=64,
        num_res_blocks=2,
        model_path=None,
        weight_dict={},
        world_size=1,
        version=2,
    ):
        super().__init__()

        self.frame_len = 17
        self.latent_len = 3 if version == 2 else 5

        base_group_norm.spatial = True if version == 2 else False

        self.encoder = VideoEncoder(
            in_channels=in_channels,
            z_channels=z_channels,
            num_res_blocks=num_res_blocks,
            version=version,
        )

        self.decoder = VideoDecoder(
            z_channels=z_channels,
            out_channels=out_channels,
            num_res_blocks=num_res_blocks,
            version=version,
        )

        if model_path is not None:
            weight_dict = self.init_from_ckpt(model_path)
        if len(weight_dict) != 0:
            self.load_from_dict(weight_dict)
        self.convert_channel_last()

        self.world_size = world_size

    def init_from_ckpt(self, model_path):
        from safetensors import safe_open