Upload batch 65: 500 files (0.01 GiB)

74da989 verified about 1 month ago

8.25 kB

from ..tools import GlobalState, DistController

from .plugins import (
    torch,
    ModulePlugin,
    GroupNormPlugin,
    Conv3DSafeNewPligin,
    Conv2DSafeNewPligin,
    WanAttentionPlugin,
    Conv2DSafeNewPliginStride2,
)
from diffusers.models.autoencoders.autoencoder_kl_wan import (
    WanCausalConv3d,
    WanAttentionBlock,
)

# from ...modules.vae import CausalConv3d, AttentionBlock
# from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d, WanAttentionBlock


class DistWrapper(object):
    def __init__(
        self, pipe, dist_controller: DistController, config, chunk_dim
    ) -> None:
        super().__init__()
        self.pipe = pipe
        self.dist_controller = dist_controller
        self.config = config
        # Ensure chunk_dim is either 3 (temporal) or 4 (spatial) for proper parallel processing
        assert chunk_dim == 3 or chunk_dim == 4
        # Initialize global state to share distributed configuration across plugins
        self.global_state = GlobalState(
            {"dist_controller": dist_controller, "chunk_dim": chunk_dim}
        )
        self.plugin_mount()

        # Configure plugin-specific hyperparameters for attention and convolution layers
        plugin_configs = {
            "attn": {
                "padding": 24,
                "top_k": 24,
                "top_k_chunk_size": 24,
                "attn_scale": 1.0,
                "token_num_scale": True,
                "dynamic_scale": True,
            },
            "conv_3d": {
                "padding": 1,
            },
            "conv_layer": {},
        }
        self.global_state.set("plugin_configs", plugin_configs)

        # torch.compile
        # torch._dynamo.config.recompile_limit = 20
        # Compile decoder with torch.compile for improved inference performance
        with torch.no_grad():
            # self.pipe.encoder = torch.compile(self.pipe.encoder)
            self.pipe.decoder = torch.compile(self.pipe.decoder)

    def plugin_mount(self):
        self.plugins = {}
        self.group_norm_plugin_mount()
        self.conv_3d_plugin_mount()
        self.conv_2d_plugin_stride2_mount()  ##only for wan vae encoder
        self.conv_2d_plugin_mount()
        self.wanattention_plugin_mount()

    def wanattention_plugin_mount(self):
        # Mount attention plugins to WanAttentionBlock modules for distributed processing
        self.plugins["wanattention"] = {}
        wanattention_s = []
        for module in self.pipe.encoder.named_modules():
            # if self.dist_controller.is_master and module[1].__class__.__name__ == 'AttentionBlock':
            #    print("Encoder attn: ", module[0])
            if (
                "mid_block." in module[0]
                and module[1].__class__.__name__ == "WanAttentionBlock"
            ):
                wanattention_s.append(module[1])
        for module in self.pipe.decoder.named_modules():
            # if self.dist_controller.is_master and module[1].__class__.__name__ == 'AttentionBlock':
            #    print("Decoder attn: ", module[0])
            if (
                "mid_block." in module[0]
                and module[1].__class__.__name__ == "WanAttentionBlock"
            ):
                wanattention_s.append(module[1])
        if self.dist_controller.is_master:
            print(f"Found {len(wanattention_s)} wanattention_s")
        for i, wanattention in enumerate(wanattention_s):
            plugin_id = "wanattention", i
            self.plugins["wanattention"][plugin_id] = WanAttentionPlugin(
                wanattention, plugin_id, self.global_state
            )

    def group_norm_plugin_mount(self):
        self.plugins["group_norm"] = {}
        group_norms = []
        for module in self.pipe.decoder.named_modules():
            if ("norm_layer" in module[0]) and module[
                1
            ].__class__.__name__ == "GroupNorm":
                group_norms.append(module[1])
        if self.dist_controller.is_master:
            print(f"Found {len(group_norms)} group norms")
        for i, group_norm in enumerate(group_norms):
            plugin_id = "group_norm", i
            self.plugins["group_norm"][plugin_id] = GroupNormPlugin(
                group_norm, plugin_id, self.global_state
            )

    def conv_3d_plugin_mount(self):
        self.plugins["conv_3d"] = {}
        conv3d_s = []
        for module in self.pipe.encoder.named_modules():
            # if isinstance(module[1], CausalConv3d):
            #    print("Encoder conv3d: ", module[0], module[1].kernel_size[1])
            if isinstance(module[1], WanCausalConv3d) and module[1].kernel_size[1] > 1:
                # print(f"Found conv3d: {module[1]}")
                conv3d_s.append(module[1])
        for module in self.pipe.decoder.named_modules():
            # if isinstance(module[1], CausalConv3d):
            #    print("Decoder conv3d: ", module[0], module[1].kernel_size[1])
            if isinstance(module[1], WanCausalConv3d) and module[1].kernel_size[1] > 1:
                # print(f"Found conv3d: {module[1]}")
                conv3d_s.append(module[1])
        if self.dist_controller.is_master:
            print(f"Found {len(conv3d_s)} conv3d_s")
        for i, conv in enumerate(conv3d_s):
            plugin_id = "conv_3d", i
            self.plugins["conv_3d"][plugin_id] = Conv3DSafeNewPligin(
                conv, plugin_id, self.global_state
            )

    def conv_2d_plugin_stride2_mount(self):
        self.plugins["conv_2d_stride2"] = {}
        conv2d_stride2_s = []
        for module in self.pipe.encoder.named_modules():
            if ".resample" in module[0] and module[1].__class__.__name__ == "Conv2d":
                conv2d_stride2_s.append(module[1])
        if self.dist_controller.is_master:
            print(f"Found {len(conv2d_stride2_s)} conv2d_stride2_s")
        for i, conv in enumerate(conv2d_stride2_s):
            plugin_id = "conv_2d_stride2", i
            self.plugins["conv_2d_stride2"][plugin_id] = Conv2DSafeNewPliginStride2(
                conv, plugin_id, self.global_state
            )

    def conv_2d_plugin_mount(self):
        self.plugins["conv_2d"] = {}
        conv2d_s = []
        for module in self.pipe.decoder.named_modules():
            if ".resample" in module[0] and module[1].__class__.__name__ == "Conv2d":
                conv2d_s.append(module[1])
        if self.dist_controller.is_master:
            print(f"Found {len(conv2d_s)} conv2d_s")
        for i, conv in enumerate(conv2d_s):
            plugin_id = "conv_2d", i
            self.plugins["conv_2d"][plugin_id] = Conv2DSafeNewPligin(
                conv, plugin_id, self.global_state
            )

    def inference(
        self,
        local_pose_image,
        local_latents,
        config={},
        pipe_configs={
            "steps": 50,
            "guidance_scale": 12,
            "fps": 60,
            "num_frames": 24 * 1,
            "height": 320,
            "width": 512,
            "export_fps": 12,
            "base_path": "./work/output",
            "file_name": None,
        },
        plugin_configs={
            "attn": {
                "padding": 24,
                "top_k": 24,
                "top_k_chunk_size": 24,
                "attn_scale": 1.0,
                "token_num_scale": True,
                "dynamic_scale": True,
            },
            "conv_3d": {
                "padding": 1,
            },
            "conv_layer": {},
        },
        additional_info={},
    ):
        # Run VAE encode-decode inference on local pose images and latents
        self.plugin_mount()
        # print("self.config seed: ", self.config["seed"])

        self.global_state.set("plugin_configs", plugin_configs)
        self.pipe = self.pipe.to(device="cuda", dtype=torch.bfloat16)
        with torch.no_grad():
            local_pose_image = local_pose_image.to(device="cuda", dtype=torch.bfloat16)
            local_latents = local_latents.to(device="cuda", dtype=torch.bfloat16)

            tmp_latents = self.pipe.encode(local_pose_image).latent_dist.mode()

            latents = self.pipe.decode(local_latents, return_dict=False)[0]
            return latents

	from ..tools import GlobalState, DistController

	from .plugins import (
	torch,
	ModulePlugin,
	GroupNormPlugin,
	Conv3DSafeNewPligin,
	Conv2DSafeNewPligin,
	WanAttentionPlugin,
	Conv2DSafeNewPliginStride2,
	)
	from diffusers.models.autoencoders.autoencoder_kl_wan import (
	WanCausalConv3d,
	WanAttentionBlock,
	)

	# from ...modules.vae import CausalConv3d, AttentionBlock
	# from diffusers.models.autoencoders.autoencoder_kl_wan import WanCausalConv3d, WanAttentionBlock


	class DistWrapper(object):
	def __init__(
	self, pipe, dist_controller: DistController, config, chunk_dim
	) -> None:
	super().__init__()
	self.pipe = pipe
	self.dist_controller = dist_controller
	self.config = config
	# Ensure chunk_dim is either 3 (temporal) or 4 (spatial) for proper parallel processing
	assert chunk_dim == 3 or chunk_dim == 4
	# Initialize global state to share distributed configuration across plugins
	self.global_state = GlobalState(
	{"dist_controller": dist_controller, "chunk_dim": chunk_dim}
	)
	self.plugin_mount()

	# Configure plugin-specific hyperparameters for attention and convolution layers
	plugin_configs = {
	"attn": {
	"padding": 24,
	"top_k": 24,
	"top_k_chunk_size": 24,
	"attn_scale": 1.0,
	"token_num_scale": True,
	"dynamic_scale": True,
	},
	"conv_3d": {
	"padding": 1,
	},
	"conv_layer": {},
	}
	self.global_state.set("plugin_configs", plugin_configs)

	# torch.compile
	# torch._dynamo.config.recompile_limit = 20
	# Compile decoder with torch.compile for improved inference performance
	with torch.no_grad():
	# self.pipe.encoder = torch.compile(self.pipe.encoder)
	self.pipe.decoder = torch.compile(self.pipe.decoder)

	def plugin_mount(self):
	self.plugins = {}
	self.group_norm_plugin_mount()
	self.conv_3d_plugin_mount()
	self.conv_2d_plugin_stride2_mount() ##only for wan vae encoder
	self.conv_2d_plugin_mount()
	self.wanattention_plugin_mount()

	def wanattention_plugin_mount(self):
	# Mount attention plugins to WanAttentionBlock modules for distributed processing
	self.plugins["wanattention"] = {}
	wanattention_s = []
	for module in self.pipe.encoder.named_modules():
	# if self.dist_controller.is_master and module[1].__class__.__name__ == 'AttentionBlock':
	# print("Encoder attn: ", module[0])
	if (
	"mid_block." in module[0]
	and module[1].__class__.__name__ == "WanAttentionBlock"
	):
	wanattention_s.append(module[1])
	for module in self.pipe.decoder.named_modules():
	# if self.dist_controller.is_master and module[1].__class__.__name__ == 'AttentionBlock':
	# print("Decoder attn: ", module[0])
	if (
	"mid_block." in module[0]
	and module[1].__class__.__name__ == "WanAttentionBlock"
	):
	wanattention_s.append(module[1])
	if self.dist_controller.is_master:
	print(f"Found {len(wanattention_s)} wanattention_s")
	for i, wanattention in enumerate(wanattention_s):
	plugin_id = "wanattention", i
	self.plugins["wanattention"][plugin_id] = WanAttentionPlugin(
	wanattention, plugin_id, self.global_state
	)

	def group_norm_plugin_mount(self):
	self.plugins["group_norm"] = {}
	group_norms = []
	for module in self.pipe.decoder.named_modules():
	if ("norm_layer" in module[0]) and module[
	1
	].__class__.__name__ == "GroupNorm":
	group_norms.append(module[1])
	if self.dist_controller.is_master:
	print(f"Found {len(group_norms)} group norms")
	for i, group_norm in enumerate(group_norms):
	plugin_id = "group_norm", i
	self.plugins["group_norm"][plugin_id] = GroupNormPlugin(
	group_norm, plugin_id, self.global_state
	)

	def conv_3d_plugin_mount(self):
	self.plugins["conv_3d"] = {}
	conv3d_s = []
	for module in self.pipe.encoder.named_modules():
	# if isinstance(module[1], CausalConv3d):
	# print("Encoder conv3d: ", module[0], module[1].kernel_size[1])
	if isinstance(module[1], WanCausalConv3d) and module[1].kernel_size[1] > 1:
	# print(f"Found conv3d: {module[1]}")
	conv3d_s.append(module[1])
	for module in self.pipe.decoder.named_modules():
	# if isinstance(module[1], CausalConv3d):
	# print("Decoder conv3d: ", module[0], module[1].kernel_size[1])
	if isinstance(module[1], WanCausalConv3d) and module[1].kernel_size[1] > 1:
	# print(f"Found conv3d: {module[1]}")
	conv3d_s.append(module[1])
	if self.dist_controller.is_master:
	print(f"Found {len(conv3d_s)} conv3d_s")
	for i, conv in enumerate(conv3d_s):
	plugin_id = "conv_3d", i
	self.plugins["conv_3d"][plugin_id] = Conv3DSafeNewPligin(
	conv, plugin_id, self.global_state
	)

	def conv_2d_plugin_stride2_mount(self):
	self.plugins["conv_2d_stride2"] = {}
	conv2d_stride2_s = []
	for module in self.pipe.encoder.named_modules():
	if ".resample" in module[0] and module[1].__class__.__name__ == "Conv2d":
	conv2d_stride2_s.append(module[1])
	if self.dist_controller.is_master:
	print(f"Found {len(conv2d_stride2_s)} conv2d_stride2_s")
	for i, conv in enumerate(conv2d_stride2_s):
	plugin_id = "conv_2d_stride2", i
	self.plugins["conv_2d_stride2"][plugin_id] = Conv2DSafeNewPliginStride2(
	conv, plugin_id, self.global_state
	)

	def conv_2d_plugin_mount(self):
	self.plugins["conv_2d"] = {}
	conv2d_s = []
	for module in self.pipe.decoder.named_modules():
	if ".resample" in module[0] and module[1].__class__.__name__ == "Conv2d":
	conv2d_s.append(module[1])
	if self.dist_controller.is_master:
	print(f"Found {len(conv2d_s)} conv2d_s")
	for i, conv in enumerate(conv2d_s):
	plugin_id = "conv_2d", i
	self.plugins["conv_2d"][plugin_id] = Conv2DSafeNewPligin(
	conv, plugin_id, self.global_state
	)

	def inference(
	self,
	local_pose_image,
	local_latents,
	config={},
	pipe_configs={
	"steps": 50,
	"guidance_scale": 12,
	"fps": 60,
	"num_frames": 24 * 1,
	"height": 320,
	"width": 512,
	"export_fps": 12,
	"base_path": "./work/output",
	"file_name": None,
	},
	plugin_configs={
	"attn": {
	"padding": 24,
	"top_k": 24,
	"top_k_chunk_size": 24,
	"attn_scale": 1.0,
	"token_num_scale": True,
	"dynamic_scale": True,
	},
	"conv_3d": {
	"padding": 1,
	},
	"conv_layer": {},
	},
	additional_info={},
	):
	# Run VAE encode-decode inference on local pose images and latents
	self.plugin_mount()
	# print("self.config seed: ", self.config["seed"])

	self.global_state.set("plugin_configs", plugin_configs)
	self.pipe = self.pipe.to(device="cuda", dtype=torch.bfloat16)
	with torch.no_grad():
	local_pose_image = local_pose_image.to(device="cuda", dtype=torch.bfloat16)
	local_latents = local_latents.to(device="cuda", dtype=torch.bfloat16)

	tmp_latents = self.pipe.encode(local_pose_image).latent_dist.mode()

	latents = self.pipe.decode(local_latents, return_dict=False)[0]
	return latents