"""
Base executor class for SGLang Diffusion ComfyUI integration.
"""

import torch


class SGLDiffusionExecutor(torch.nn.Module):
    """Base executor class for SGLang Diffusion models in ComfyUI."""

    def __init__(self, generator, model_path, model, config):
        super(SGLDiffusionExecutor, self).__init__()
        self.generator = generator
        self.model_path = model_path
        self.model = model
        self.dtype = config.unet_config["dtype"]
        self.config = config
        self.loras = []

    def _unpack_latents(self, latents, height, width, channels):
        """Unpack latents from packed format to standard format."""
        batch_size = latents.shape[0]
        latents = latents.view(batch_size, height // 2, width // 2, channels, 2, 2)
        latents = latents.permute(0, 3, 1, 4, 2, 5)
        latents = latents.reshape(batch_size, channels, height, width)

        return latents

    def _pack_latents(self, latents):
        """Pack latents from standard format to packed format."""
        batch_size, num_channels_latents, height, width = latents.shape
        latents = latents.view(
            batch_size, num_channels_latents, height // 2, 2, width // 2, 2
        )
        latents = latents.permute(0, 2, 4, 1, 3, 5)
        latents = latents.reshape(
            batch_size, (height // 2) * (width // 2), num_channels_latents * 4
        )
        return latents
