# Path: third_party/models/networks.py
# def init_net(net, gpu_ids=[]):
#     """Initialize a network by registering CPU/GPU device (with multi-GPU support)
#     Parameters:
#         net (network)      -- the network to be initialized
#         gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
# 
#     Return an initialized network.
#     """
#     if len(gpu_ids) > 0:
#         assert (torch.cuda.is_available())
#         net.to(gpu_ids[0])
#         net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs
#     return net


import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from third_party.models.networks import init_net


###############################################################################
# Helper Functions
###############################################################################
def define_LNR(nf=64, texture_channels=16, texture_res=16, n_textures=25, gpu_ids=[]):
    """Create a layered neural renderer.

    Parameters:
        nf (int) -- the number of channels in the first/last conv layers
        texture_channels (int) -- the number of channels in the neural texture
        texture_res (int) -- the size of each individual texture map
        n_textures (int) -- the number of individual texture maps
        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2

    Returns a layered neural rendering model.
    """
    net = LayeredNeuralRenderer(nf, texture_channels, texture_res, n_textures)
    return init_net(net, gpu_ids)


def define_kp2uv(nf=64, gpu_ids=[]):
    """Create a keypoint-to-UV model.

    Parameters:
        nf (int) -- the number of channels in the first/last conv layers

    Returns a keypoint-to-UV model.
    """
    net = kp2uv(nf)

# Path: third_party/models/networks.py
# def init_net(net, gpu_ids=[]):
#     """Initialize a network by registering CPU/GPU device (with multi-GPU support)
#     Parameters:
#         net (network)      -- the network to be initialized
#         gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2
# 
#     Return an initialized network.
#     """
#     if len(gpu_ids) > 0:
#         assert (torch.cuda.is_available())
#         net.to(gpu_ids[0])
#         net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs
#     return net


# Path: models/networks.py
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from third_party.models.networks import init_net



###############################################################################
# Helper Functions
###############################################################################
def define_LNR(nf=64, texture_channels=16, texture_res=16, n_textures=25, gpu_ids=[]):
    """Create a layered neural renderer.

    Parameters:
        nf (int) -- the number of channels in the first/last conv layers
        texture_channels (int) -- the number of channels in the neural texture
        texture_res (int) -- the size of each individual texture map
        n_textures (int) -- the number of individual texture maps
        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2

    Returns a layered neural rendering model.
    """
    net = LayeredNeuralRenderer(nf, texture_channels, texture_res, n_textures)
    return init_net(net, gpu_ids)


def define_kp2uv(nf=64, gpu_ids=[]):
    """Create a keypoint-to-UV model.

    Parameters:
        nf (int) -- the number of channels in the first/last conv layers

    Returns a keypoint-to-UV model.
    """
    net = kp2uv(nf)

    return init_net(net, gpu_ids)