Function forward_G

draggan/draggan.py:116–140

Forward pass through the generator network. Args: G (torch.nn.Module): The generator network. W (torch.Tensor): The latent code tensor of shape [batch_size, latent_dim, 512]. device (torch.device): The device to use for the computation. Returns: A tuple

(
    G: torch.nn.Module,
    W: torch.Tensor,
    device: torch.device,
)

def forward_G(
    G: torch.nn.Module,
    W: torch.Tensor,
    device: torch.device,
) -> Tuple[torch.Tensor, torch.Tensor]:
    """
    Forward pass through the generator network.

    Args:
        G (torch.nn.Module): The generator network.
        W (torch.Tensor): The latent code tensor of shape [batch_size, latent_dim, 512].
        device (torch.device): The device to use for the computation.

    Returns:
        A tuple containing the generated image tensor of shape [batch_size, 3, height, width]
        and the feature maps tensor of shape [batch_size, num_channels, height, width].
    """
    register_hook(G)
    
    if not isinstance(W, torch.Tensor):
        W = torch.from_numpy(W).to(device)

    img = G.synthesis(W, noise_mode="const", force_fp32=True)

    return img, G.activations[0]


def generate_image(