# Path: nn_utils/activations.py
# def to_hdr_activation(x):
#     return tf.exp(tf.nn.relu(x)) - 1
# 
# Path: nn_utils/film_siren_layers.py
# class FiLMSiren(tf.keras.layers.Layer):
#     def __init__(self, w0=30.0, **kwargs):
#         super(FiLMSiren, self).__init__(trainable=False, **kwargs)
#         self.sine = Sine(w0=w0)
# 
#     @tf.function
#     def call(self, tensor, frequency, phase_shift):
#         return self.sine(tensor * frequency + phase_shift)
# 
# Path: nn_utils/initializers.py
# class SIRENFirstLayerInitializer(tf.keras.initializers.Initializer):
#     def __init__(self, scale=1.0, seed=None):
#         super().__init__()
#         self.scale = scale
#         self.seed = seed
# 
#     def __call__(self, shape, dtype=tf.float32):
#         fan_in, _ = _compute_fans(shape)
#         limit = self.scale / max(1.0, float(fan_in))
#         return tf.random.uniform(shape, -limit, limit, seed=self.seed, dtype=dtype)
# 
#     def get_config(self):
#         base_config = super().get_config()
#         config = {"scale": self.scale, "seed": self.seed}
#         return dict(list(base_config.items()) + list(config.items()))
# 
# class SIRENInitializer(tf.keras.initializers.Initializer):
#     def __init__(self, w0: float = 30.0, c: float = 6.0, seed: int = None):
#         super().__init__()
#         self.w0 = w0
#         self.c = c
#         self.seed = seed
# 
#     def __call__(self, shape, dtype=tf.float32):
#         fan_in, _ = _compute_fans(shape)
#         limit = tf.math.sqrt(self.c / max(1.0, float(fan_in))) / self.w0
#         return tf.random.uniform(shape, -limit, limit, seed=self.seed, dtype=dtype)
# 
#     def get_config(self):
#         base_config = super().get_config()
#         config = {"w0": self.w0, "c": self.c, "seed": self.seed}
#         return dict(list(base_config.items()) + list(config.items()))
# 
# Path: nn_utils/nerf_layers.py
# class FourierEmbedding(tf.keras.layers.Layer):
#     def __init__(
#         self,
#         num_frequencies: int,
#         include_input: bool = True,
#         input_dim: int = 3,
#         **kwargs
#     ):
#         super(FourierEmbedding, self).__init__(trainable=False, **kwargs)
# 
#         self.input_dims = input_dim
#         self.out_dims = 0
#         self.num_frequencies = num_frequencies
#         self.include_input = include_input
# 
#         if include_input:
#             self.out_dims += input_dim
# 
#         if num_frequencies >= 1:
#             frequency_dims = (
#                 2 * num_frequencies * input_dim
#             )  # 2 (sin, cos) * num_frequencies * num input_dim
#             self.out_dims += frequency_dims
# 
#         self.scales = tf.convert_to_tensor(
#             2.0 ** np.linspace(0.0, self.num_frequencies - 1, self.num_frequencies),
#             tf.float32,
#         )
# 
#     # @tf.function
#     def call(self, x):
#         assert (
#             x.shape[-1] == self.input_dims
#         ), "Channel dimension is %d but should be %d" % (x.shape[-1], self.input_dims)
# 
#         x_shape = tf.shape(x)
# 
#         xb = tf.reshape(
#             x[..., None, :] * self.scales[:, None],
#             tf.concat([x_shape[:-1], self.num_frequencies * x_shape[-1:]], 0),
#         )
# 
#         four_feat = tf.math.sin(tf.concat([xb, xb + 0.5 * np.pi], axis=-1))
# 
#         ret = []
#         if self.include_input:
#             ret += [x]
# 
#         return tf.concat(ret + [four_feat], axis=-1)
# 
#     def get_output_dimensionality(self):
#         return self.out_dims


import tensorflow as tf
import tensorflow.keras as K
import tensorflow.keras.layers as layers
import nn_utils.math_utils as math_utils
        import numpy as np
from nn_utils.activations import to_hdr_activation
from nn_utils.film_siren_layers import FiLMSiren
from nn_utils.initializers import SIRENFirstLayerInitializer, SIRENInitializer
from nn_utils.nerf_layers import FourierEmbedding
        self.film_layer = FiLMSiren(w0=1)

        self.net = []
        input_size = 3
        for i in range(args.net_depth):
            self.net.append(
                K.Sequential(
                    [
                        layers.InputLayer(
                            (input_size,), name="MainNetwork/Input%d" % i
                        ),
                        layers.Dense(
                            args.net_width,
                            activation=None,
                            kernel_initializer=SIRENFirstLayerInitializer(0.5)
                            if i == 0
                            else SIRENInitializer(w0=2),
                            name="MainNetwork/Layer%d" % i,
                        ),
                    ]
                )
            )
            input_size = args.net_width

        self.conditional_layer = K.Sequential(
            [
                layers.InputLayer((input_size,), name="MainNetwork/ConditionalInput"),
                layers.Dense(
                    args.net_width,
                    activation=None,

# Path: nn_utils/activations.py
# def to_hdr_activation(x):
#     return tf.exp(tf.nn.relu(x)) - 1
# 
# Path: nn_utils/film_siren_layers.py
# class FiLMSiren(tf.keras.layers.Layer):
#     def __init__(self, w0=30.0, **kwargs):
#         super(FiLMSiren, self).__init__(trainable=False, **kwargs)
#         self.sine = Sine(w0=w0)
# 
#     @tf.function
#     def call(self, tensor, frequency, phase_shift):
#         return self.sine(tensor * frequency + phase_shift)
# 
# Path: nn_utils/initializers.py
# class SIRENFirstLayerInitializer(tf.keras.initializers.Initializer):
#     def __init__(self, scale=1.0, seed=None):
#         super().__init__()
#         self.scale = scale
#         self.seed = seed
# 
#     def __call__(self, shape, dtype=tf.float32):
#         fan_in, _ = _compute_fans(shape)
#         limit = self.scale / max(1.0, float(fan_in))
#         return tf.random.uniform(shape, -limit, limit, seed=self.seed, dtype=dtype)
# 
#     def get_config(self):
#         base_config = super().get_config()
#         config = {"scale": self.scale, "seed": self.seed}
#         return dict(list(base_config.items()) + list(config.items()))
# 
# class SIRENInitializer(tf.keras.initializers.Initializer):
#     def __init__(self, w0: float = 30.0, c: float = 6.0, seed: int = None):
#         super().__init__()
#         self.w0 = w0
#         self.c = c
#         self.seed = seed
# 
#     def __call__(self, shape, dtype=tf.float32):
#         fan_in, _ = _compute_fans(shape)
#         limit = tf.math.sqrt(self.c / max(1.0, float(fan_in))) / self.w0
#         return tf.random.uniform(shape, -limit, limit, seed=self.seed, dtype=dtype)
# 
#     def get_config(self):
#         base_config = super().get_config()
#         config = {"w0": self.w0, "c": self.c, "seed": self.seed}
#         return dict(list(base_config.items()) + list(config.items()))
# 
# Path: nn_utils/nerf_layers.py
# class FourierEmbedding(tf.keras.layers.Layer):
#     def __init__(
#         self,
#         num_frequencies: int,
#         include_input: bool = True,
#         input_dim: int = 3,
#         **kwargs
#     ):
#         super(FourierEmbedding, self).__init__(trainable=False, **kwargs)
# 
#         self.input_dims = input_dim
#         self.out_dims = 0
#         self.num_frequencies = num_frequencies
#         self.include_input = include_input
# 
#         if include_input:
#             self.out_dims += input_dim
# 
#         if num_frequencies >= 1:
#             frequency_dims = (
#                 2 * num_frequencies * input_dim
#             )  # 2 (sin, cos) * num_frequencies * num input_dim
#             self.out_dims += frequency_dims
# 
#         self.scales = tf.convert_to_tensor(
#             2.0 ** np.linspace(0.0, self.num_frequencies - 1, self.num_frequencies),
#             tf.float32,
#         )
# 
#     # @tf.function
#     def call(self, x):
#         assert (
#             x.shape[-1] == self.input_dims
#         ), "Channel dimension is %d but should be %d" % (x.shape[-1], self.input_dims)
# 
#         x_shape = tf.shape(x)
# 
#         xb = tf.reshape(
#             x[..., None, :] * self.scales[:, None],
#             tf.concat([x_shape[:-1], self.num_frequencies * x_shape[-1:]], 0),
#         )
# 
#         four_feat = tf.math.sin(tf.concat([xb, xb + 0.5 * np.pi], axis=-1))
# 
#         ret = []
#         if self.include_input:
#             ret += [x]
# 
#         return tf.concat(ret + [four_feat], axis=-1)
# 
#     def get_output_dimensionality(self):
#         return self.out_dims


# Path: models/illumination_integration_net/model.py
import tensorflow as tf
import tensorflow.keras as K
import tensorflow.keras.layers as layers
import nn_utils.math_utils as math_utils
        import numpy as np
from nn_utils.activations import to_hdr_activation
from nn_utils.film_siren_layers import FiLMSiren
from nn_utils.initializers import SIRENFirstLayerInitializer, SIRENInitializer
from nn_utils.nerf_layers import FourierEmbedding

        self.film_layer = FiLMSiren(w0=1)

        self.net = []
        input_size = 3
        for i in range(args.net_depth):
            self.net.append(
                K.Sequential(
                    [
                        layers.InputLayer(
                            (input_size,), name="MainNetwork/Input%d" % i
                        ),
                        layers.Dense(
                            args.net_width,
                            activation=None,
                            kernel_initializer=SIRENFirstLayerInitializer(0.5)
                            if i == 0
                            else SIRENInitializer(w0=2),
                            name="MainNetwork/Layer%d" % i,
                        ),
                    ]
                )
            )
            input_size = args.net_width

        self.conditional_layer = K.Sequential(
            [
                layers.InputLayer((input_size,), name="MainNetwork/ConditionalInput"),
                layers.Dense(
                    args.net_width,
                    activation=None,

                    kernel_initializer=SIRENInitializer(w0=2),