Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
model.py 5 KB
Newer Older
helloyongyang's avatar
first commit  
helloyongyang committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133 
import os
import torch
import time
import glob
from lightx2v.text2v.models.networks.wan.weights.pre_weights import WanPreWeights
from lightx2v.text2v.models.networks.wan.weights.post_weights import WanPostWeights
from lightx2v.text2v.models.networks.wan.weights.transformer_weights import (
    WanTransformerWeights,
)
from lightx2v.text2v.models.networks.wan.infer.pre_infer import WanPreInfer
from lightx2v.text2v.models.networks.wan.infer.post_infer import WanPostInfer
from lightx2v.text2v.models.networks.wan.infer.transformer_infer import (
    WanTransformerInfer,
)
from lightx2v.text2v.models.networks.wan.infer.feature_caching.transformer_infer import WanTransformerInferFeatureCaching
from safetensors import safe_open


class WanModel:
    pre_weight_class = WanPreWeights
    post_weight_class = WanPostWeights
    transformer_weight_class = WanTransformerWeights

    def __init__(self, model_path, config):
        self.model_path = model_path
        self.config = config
        self._init_infer_class()
        self._init_weights()
        self._init_infer()

    def _init_infer_class(self):
        self.pre_infer_class = WanPreInfer
        self.post_infer_class = WanPostInfer
        if self.config["feature_caching"] == "NoCaching":
            self.transformer_infer_class = WanTransformerInfer
        elif self.config["feature_caching"] == "Tea":
            self.transformer_infer_class = WanTransformerInferFeatureCaching
        else:
            raise NotImplementedError(
                f"Unsupported feature_caching type: {self.config['feature_caching']}"
            )

    def _load_safetensor_to_dict(self, file_path):
        with safe_open(file_path, framework="pt") as f:
            tensor_dict = {
                key: f.get_tensor(key).to(torch.bfloat16).cuda() for key in f.keys()
            }
        return tensor_dict

    def _load_ckpt(self):
        safetensors_pattern = os.path.join(self.model_path, "*.safetensors")
        safetensors_files = glob.glob(safetensors_pattern)

        if not safetensors_files:
            raise FileNotFoundError(
                f"No .safetensors files found in directory: {self.model_path}"
            )
        weight_dict = {}
        for file_path in safetensors_files:
            file_weights = self._load_safetensor_to_dict(file_path)
            weight_dict.update(file_weights)
        return weight_dict

    def _init_weights(self):
        weight_dict = self._load_ckpt()
        # init weights
        self.pre_weight = self.pre_weight_class(self.config)
        self.post_weight = self.post_weight_class()
        self.transformer_weights = self.transformer_weight_class(self.config)
        # load weights
        self.pre_weight.load_weights(weight_dict)
        self.post_weight.load_weights(weight_dict)
        self.transformer_weights.load_weights(weight_dict)

    def _init_infer(self):
        self.pre_infer = self.pre_infer_class(self.config)
        self.post_infer = self.post_infer_class(self.config)
        self.transformer_infer = self.transformer_infer_class(self.config)

    def set_scheduler(self, scheduler):
        self.scheduler = scheduler
        self.transformer_infer.set_scheduler(scheduler)

    @torch.no_grad()
    def infer(self, text_encoders_output, image_encoder_output, args):

        timestep = torch.stack([self.scheduler.timesteps[self.scheduler.step_index]])

        embed, grid_sizes, pre_infer_out = self.pre_infer.infer(
            self.pre_weight,
            [self.scheduler.latents],
            timestep,
            text_encoders_output["context"],
            self.scheduler.seq_len,
            image_encoder_output["clip_encoder_out"],
            [image_encoder_output["vae_encode_out"]],
        )
        x = self.transformer_infer.infer(
            self.transformer_weights, grid_sizes, embed, *pre_infer_out
        )
        noise_pred_cond = self.post_infer.infer(
            self.post_weight, x, embed, grid_sizes
        )[0]

        if self.config["feature_caching"] == "Tea":
            self.scheduler.cnt += 1
            if self.scheduler.cnt >= self.scheduler.num_steps:
                self.scheduler.cnt = 0

        embed, grid_sizes, pre_infer_out = self.pre_infer.infer(
            self.pre_weight,
            [self.scheduler.latents],
            timestep,
            text_encoders_output["context_null"],
            self.scheduler.seq_len,
            image_encoder_output["clip_encoder_out"],
            [image_encoder_output["vae_encode_out"]],
        )
        x = self.transformer_infer.infer(
            self.transformer_weights, grid_sizes, embed, *pre_infer_out
        )
        noise_pred_uncond = self.post_infer.infer(
            self.post_weight, x, embed, grid_sizes
        )[0]

        if self.config["feature_caching"] == "Tea":
            self.scheduler.cnt += 1
            if self.scheduler.cnt >= self.scheduler.num_steps:
                self.scheduler.cnt = 0

        self.scheduler.noise_pred = noise_pred_uncond + args.sample_guide_scale * (
            noise_pred_cond - noise_pred_uncond
        )