Attention

        act_layer=nn.GELU,
        bias=True,
        drop=0.0,
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        bias = to_2tuple(bias)
        drop_probs = to_2tuple(drop)

        self.fc1 = nn.Linear(in_features, hidden_features, bias=bias[0])
        self.act = act_layer()
        self.drop1 = nn.Dropout(drop_probs[0])
        self.fc2 = nn.Linear(hidden_features, out_features, bias=bias[1])
        self.drop2 = nn.Dropout(drop_probs[1])

    def forward(self, x):
        return self.drop2(self.fc2(self.drop1(self.act(self.fc1(x)))))

class Attention(nn.Module):

    def __init__(
        self, dim, rope=None, num_heads=8, qkv_bias=False, attn_drop=0.0, proj_drop=0.0
    ):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim**-0.5
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)
        self.rope = rope.float() if rope is not None else None

    def forward(self, x, xpos):
        B, N, C = x.shape

        qkv = (
            self.qkv(x)
            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
            .transpose(1, 3)
        )
        q, k, v = [qkv[:, :, i] for i in range(3)]

        q_type = q.dtype
        k_type = k.dtype
        if self.rope is not None:
            q = q.float()
            k = k.float()
            with torch.autocast(device_type="cuda", enabled=False):
                q = self.rope(q, xpos)
                k = self.rope(k, xpos)
            q = q.to(q_type)
            k = k.to(k_type)

        x = (
            scaled_dot_product_attention(
                query=q, key=k, value=v, dropout_p=self.attn_drop.p, scale=self.scale
            )
            .transpose(1, 2)
            .reshape(B, N, C)
        )

        x = self.proj(x)
        x = self.proj_drop(x)
        return x

class Block(nn.Module):

    def __init__(
        self,
        dim,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        rope=None,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(


Block

            k = k.float()
            with torch.autocast(device_type="cuda", enabled=False):
                q = self.rope(q, xpos)
                k = self.rope(k, xpos)
            q = q.to(q_type)
            k = k.to(k_type)

        x = (
            scaled_dot_product_attention(
                query=q, key=k, value=v, dropout_p=self.attn_drop.p, scale=self.scale
            )
            .transpose(1, 2)
            .reshape(B, N, C)
        )

        x = self.proj(x)
        x = self.proj_drop(x)
        return x

class Block(nn.Module):

    def __init__(
        self,
        dim,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        rope=None,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(
            dim,
            rope=rope,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            attn_drop=attn_drop,
            proj_drop=drop,
        )

        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )

    def forward(self, x, xpos):
        x = x + self.drop_path(self.attn(self.norm1(x), xpos))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x

class CrossAttention(nn.Module):

    def __init__(
        self, dim, rope=None, num_heads=8, qkv_bias=False, attn_drop=0.0, proj_drop=0.0
    ):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim**-0.5

        self.projq = nn.Linear(dim, dim, bias=qkv_bias)
        self.projk = nn.Linear(dim, dim, bias=qkv_bias)
        self.projv = nn.Linear(dim, dim, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        self.rope = rope.float() if rope is not None else None


CrossAttention

            attn_drop=attn_drop,
            proj_drop=drop,
        )

        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )

    def forward(self, x, xpos):
        x = x + self.drop_path(self.attn(self.norm1(x), xpos))
        x = x + self.drop_path(self.mlp(self.norm2(x)))
        return x

class CrossAttention(nn.Module):

    def __init__(
        self, dim, rope=None, num_heads=8, qkv_bias=False, attn_drop=0.0, proj_drop=0.0
    ):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = head_dim**-0.5

        self.projq = nn.Linear(dim, dim, bias=qkv_bias)
        self.projk = nn.Linear(dim, dim, bias=qkv_bias)
        self.projv = nn.Linear(dim, dim, bias=qkv_bias)
        self.attn_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

        self.rope = rope.float() if rope is not None else None

    def forward(self, query, key, value, qpos, kpos, use_ttt3r=False):
        B, Nq, C = query.shape
        Nk = key.shape[1]
        Nv = value.shape[1]

        q = (
            self.projq(query)
            .reshape(B, Nq, self.num_heads, C // self.num_heads)
            .permute(0, 2, 1, 3)
        )
        k = (
            self.projk(key)
            .reshape(B, Nk, self.num_heads, C // self.num_heads)
            .permute(0, 2, 1, 3)
        )
        v = (
            self.projv(value)
            .reshape(B, Nv, self.num_heads, C // self.num_heads)
            .permute(0, 2, 1, 3)
        )

        q_type = q.dtype
        k_type = k.dtype
        if self.rope is not None:
            if qpos is not None:
                q = q.float()
                with torch.autocast(device_type="cuda", enabled=False):
                    q = self.rope(q, qpos)
                q = q.to(q_type)

            if kpos is not None:
                k = k.float()
                with torch.autocast(device_type="cuda", enabled=False):
                    k = self.rope(k, kpos)
                k = k.to(k_type)

        if use_ttt3r:
            # original attention
            attn = (q @ k.transpose(-2, -1)) * self.scale  # [B, num_heads, Nq, Nk] [1, 16, 768, 1 + 576]
            attn_before_softmax = attn.detach().clone()

            attn = attn.softmax(dim=-1)
            attn = self.attn_drop(attn)
            x = (attn @ v).transpose(1, 2).reshape(B, Nq, C) # [B, Nq, C] [1, 768, 768]

            x = self.proj(x)
            x = self.proj_drop(x)

            return x, attn_before_softmax
        else:
            x = (
                scaled_dot_product_attention(
                    query=q, key=k, value=v, dropout_p=self.attn_drop.p, scale=self.scale
                )
                .transpose(1, 2)
                .reshape(B, Nq, C)
            )

            x = self.proj(x)
            x = self.proj_drop(x)
            return x, None

class DecoderBlock(nn.Module):

    def __init__(
        self,
        dim,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        norm_mem=True,
        rope=None,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim)


DecoderBlock

            x = (attn @ v).transpose(1, 2).reshape(B, Nq, C) # [B, Nq, C] [1, 768, 768]

            x = self.proj(x)
            x = self.proj_drop(x)

            return x, attn_before_softmax
        else:
            x = (
                scaled_dot_product_attention(
                    query=q, key=k, value=v, dropout_p=self.attn_drop.p, scale=self.scale
                )
                .transpose(1, 2)
                .reshape(B, Nq, C)
            )

            x = self.proj(x)
            x = self.proj_drop(x)
            return x, None

class DecoderBlock(nn.Module):

    def __init__(
        self,
        dim,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        norm_mem=True,
        rope=None,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(
            dim,
            rope=rope,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            attn_drop=attn_drop,
            proj_drop=drop,
        )
        self.cross_attn = CrossAttention(
            dim,
            rope=rope,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            attn_drop=attn_drop,
            proj_drop=drop,
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = norm_layer(dim)
        self.norm3 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )
        self.norm_y = norm_layer(dim) if norm_mem else nn.Identity()

    def forward(self, x, y, xpos, ypos, use_ttt3r=False):
        x = x + self.drop_path(self.attn(self.norm1(x), xpos))
        y_ = self.norm_y(y)
        cross_attn_output, cross_attn = self.cross_attn(self.norm2(x), y_, y_, xpos, ypos, use_ttt3r=use_ttt3r)
        x = x + self.drop_path(cross_attn_output)
        x = x + self.drop_path(self.mlp(self.norm3(x)))
        return x, y, cross_attn

class CustomDecoderBlock(nn.Module):

    def __init__(
        self,
        dim,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        norm_mem=True,
        rope=None,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim)


CustomDecoderBlock

        self.norm2 = norm_layer(dim)
        self.norm3 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )
        self.norm_y = norm_layer(dim) if norm_mem else nn.Identity()

    def forward(self, x, y, xpos, ypos, use_ttt3r=False):
        x = x + self.drop_path(self.attn(self.norm1(x), xpos))
        y_ = self.norm_y(y)
        cross_attn_output, cross_attn = self.cross_attn(self.norm2(x), y_, y_, xpos, ypos, use_ttt3r=use_ttt3r)
        x = x + self.drop_path(cross_attn_output)
        x = x + self.drop_path(self.mlp(self.norm3(x)))
        return x, y, cross_attn

class CustomDecoderBlock(nn.Module):

    def __init__(
        self,
        dim,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        norm_mem=True,
        rope=None,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = Attention(
            dim,
            rope=rope,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            attn_drop=attn_drop,
            proj_drop=drop,
        )
        self.cross_attn = CrossAttention(
            dim,
            rope=rope,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            attn_drop=attn_drop,
            proj_drop=drop,
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = norm_layer(dim)
        self.norm3 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )
        self.norm_y = norm_layer(dim) if norm_mem else nn.Identity()
        self.norm_z = norm_layer(dim) if norm_mem else nn.Identity()

    def forward(self, x, y, z, xpos, ypos):
        x = x + self.drop_path(self.attn(self.norm1(x), xpos))
        y_ = self.norm_y(y)
        z_ = self.norm_z(z)
        x = x + self.drop_path(self.cross_attn(self.norm2(x), y_, z_, xpos, ypos))
        x = x + self.drop_path(self.mlp(self.norm3(x)))
        return x, y

class ModLN(nn.Module):
    """
    Modulation with adaLN.

    References:
    DiT: https://github.com/facebookresearch/DiT/blob/main/models.py#L101
    """

    def __init__(self, inner_dim: int, mod_dim: int, eps: float):
        super().__init__()
        self.norm = nn.LayerNorm(inner_dim, eps=eps)
        self.mlp = nn.Sequential(
            nn.SiLU(),
            nn.Linear(mod_dim, inner_dim * 2),
        )

    @staticmethod
    def modulate(x, shift, scale):


ModLN

        self.norm3 = norm_layer(dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )
        self.norm_y = norm_layer(dim) if norm_mem else nn.Identity()
        self.norm_z = norm_layer(dim) if norm_mem else nn.Identity()

    def forward(self, x, y, z, xpos, ypos):
        x = x + self.drop_path(self.attn(self.norm1(x), xpos))
        y_ = self.norm_y(y)
        z_ = self.norm_z(z)
        x = x + self.drop_path(self.cross_attn(self.norm2(x), y_, z_, xpos, ypos))
        x = x + self.drop_path(self.mlp(self.norm3(x)))
        return x, y

class ModLN(nn.Module):
    """
    Modulation with adaLN.

    References:
    DiT: https://github.com/facebookresearch/DiT/blob/main/models.py#L101
    """

    def __init__(self, inner_dim: int, mod_dim: int, eps: float):
        super().__init__()
        self.norm = nn.LayerNorm(inner_dim, eps=eps)
        self.mlp = nn.Sequential(
            nn.SiLU(),
            nn.Linear(mod_dim, inner_dim * 2),
        )

    @staticmethod
    def modulate(x, shift, scale):

        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

    def forward(self, x: torch.Tensor, mod: torch.Tensor) -> torch.Tensor:
        shift, scale = self.mlp(mod).chunk(2, dim=-1)  # [N, D]
        return self.modulate(self.norm(x), shift, scale)  # [N, L, D]

class ConditionModulationBlock(nn.Module):

    def __init__(
        self,
        dim,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=partial(ModLN, eps=1e-6),
        rope=None,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim, dim)
        self.attn = Attention(


ConditionModulationBlock

    """

    def __init__(self, inner_dim: int, mod_dim: int, eps: float):
        super().__init__()
        self.norm = nn.LayerNorm(inner_dim, eps=eps)
        self.mlp = nn.Sequential(
            nn.SiLU(),
            nn.Linear(mod_dim, inner_dim * 2),
        )

    @staticmethod
    def modulate(x, shift, scale):

        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

    def forward(self, x: torch.Tensor, mod: torch.Tensor) -> torch.Tensor:
        shift, scale = self.mlp(mod).chunk(2, dim=-1)  # [N, D]
        return self.modulate(self.norm(x), shift, scale)  # [N, L, D]

class ConditionModulationBlock(nn.Module):

    def __init__(
        self,
        dim,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=partial(ModLN, eps=1e-6),
        rope=None,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim, dim)
        self.attn = Attention(
            dim,
            rope=rope,
            num_heads=num_heads,
            qkv_bias=qkv_bias,
            attn_drop=attn_drop,
            proj_drop=drop,
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = norm_layer(dim, dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )

    def forward(self, x, mod, xpos):
        x = x + self.drop_path(self.attn(self.norm1(x, mod), xpos))
        x = x + self.drop_path(self.mlp(self.norm2(x, mod)))
        return x

class PositionGetter(object):
    """return positions of patches"""

    def __init__(self):
        self.cache_positions = {}

    def __call__(self, b, h, w, device):
        if not (h, w) in self.cache_positions:
            x = torch.arange(w, device=device)
            y = torch.arange(h, device=device)
            self.cache_positions[h, w] = torch.cartesian_prod(y, x)  # (h, w, 2)
        pos = self.cache_positions[h, w].view(1, h * w, 2).expand(b, -1, 2).clone()
        return pos

class PatchEmbed(nn.Module):
    """just adding _init_weights + position getter compared to timm.models.layers.patch_embed.PatchEmbed"""


PositionGetter

            qkv_bias=qkv_bias,
            attn_drop=attn_drop,
            proj_drop=drop,
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
        self.norm2 = norm_layer(dim, dim)
        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )

    def forward(self, x, mod, xpos):
        x = x + self.drop_path(self.attn(self.norm1(x, mod), xpos))
        x = x + self.drop_path(self.mlp(self.norm2(x, mod)))
        return x

class PositionGetter(object):
    """return positions of patches"""

    def __init__(self):
        self.cache_positions = {}

    def __call__(self, b, h, w, device):
        if not (h, w) in self.cache_positions:
            x = torch.arange(w, device=device)
            y = torch.arange(h, device=device)
            self.cache_positions[h, w] = torch.cartesian_prod(y, x)  # (h, w, 2)
        pos = self.cache_positions[h, w].view(1, h * w, 2).expand(b, -1, 2).clone()
        return pos

class PatchEmbed(nn.Module):
    """just adding _init_weights + position getter compared to timm.models.layers.patch_embed.PatchEmbed"""

    def __init__(
        self,
        img_size=224,
        patch_size=16,
        in_chans=3,
        embed_dim=768,
        norm_layer=None,
        flatten=True,
    ):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])


PatchEmbed

        x = x + self.drop_path(self.attn(self.norm1(x, mod), xpos))
        x = x + self.drop_path(self.mlp(self.norm2(x, mod)))
        return x

class PositionGetter(object):
    """return positions of patches"""

    def __init__(self):
        self.cache_positions = {}

    def __call__(self, b, h, w, device):
        if not (h, w) in self.cache_positions:
            x = torch.arange(w, device=device)
            y = torch.arange(h, device=device)
            self.cache_positions[h, w] = torch.cartesian_prod(y, x)  # (h, w, 2)
        pos = self.cache_positions[h, w].view(1, h * w, 2).expand(b, -1, 2).clone()
        return pos

class PatchEmbed(nn.Module):
    """just adding _init_weights + position getter compared to timm.models.layers.patch_embed.PatchEmbed"""

    def __init__(
        self,
        img_size=224,
        patch_size=16,
        in_chans=3,
        embed_dim=768,
        norm_layer=None,
        flatten=True,
    ):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
        self.num_patches = self.grid_size[0] * self.grid_size[1]
        self.flatten = flatten

        self.proj = nn.Conv2d(
            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
        )
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

        self.position_getter = PositionGetter()

    def forward(self, x):
        B, C, H, W = x.shape
        torch._assert(
            H == self.img_size[0],
            f"Input image height ({H}) doesn't match model ({self.img_size[0]}).",
        )
        torch._assert(
            W == self.img_size[1],
            f"Input image width ({W}) doesn't match model ({self.img_size[1]}).",
        )
        x = self.proj(x)
        pos = self.position_getter(B, x.size(2), x.size(3), x.device)
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
        x = self.norm(x)
        return x, pos

    def _init_weights(self):
        w = self.proj.weight.data
        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))

if __name__ == "__main__":
    import os
    import sys

    sys.path.append(os.path.dirname(os.path.dirname(__file__)))
    import dust3r.utils.path_to_croco
    from models.pos_embed import get_2d_sincos_pos_embed, RoPE2D
    from functools import partial
    from torch.utils.checkpoint import checkpoint

    torch.manual_seed(0)

    enc_blocks_ray_map = (
        nn.ModuleList(
            [
                Block(
                    768,
                    16,


parse

# Copyright (C) 2024-present Naver Corporation. All rights reserved.
# Licensed under CC BY-NC-SA 4.0 (non-commercial use only).
#
# --------------------------------------------------------
# modified from DUSt3R

import torch
import torch.nn as nn

from itertools import repeat
import collections.abc
from torch.nn.functional import scaled_dot_product_attention
from functools import partial

def _ntuple(n):
    def parse(x):
        if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
            return x
        return tuple(repeat(x, n))

    return parse

to_2tuple = _ntuple(2)

def drop_path(
    x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True
):
    """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
    if drop_prob == 0.0 or not training:
        return x
    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (
        x.ndim - 1
    )  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
    if keep_prob > 0.0 and scale_by_keep:
        random_tensor.div_(keep_prob)


class PositionGetter(object):
    """return positions of patches"""

    def __init__(self):
        self.cache_positions = {}

    def __call__(self, b, h, w, device):
        if not (h, w) in self.cache_positions:
            x = torch.arange(w, device=device)
            y = torch.arange(h, device=device)
            self.cache_positions[h, w] = torch.cartesian_prod(y, x)  # (h, w, 2)
        pos = self.cache_positions[h, w].view(1, h * w, 2).expand(b, -1, 2).clone()
        return pos

class PatchEmbed(nn.Module):
    """just adding _init_weights + position getter compared to timm.models.layers.patch_embed.PatchEmbed"""

    def __init__(
        self,
        img_size=224,
        patch_size=16,
        in_chans=3,
        embed_dim=768,
        norm_layer=None,
        flatten=True,
    ):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
        self.num_patches = self.grid_size[0] * self.grid_size[1]
        self.flatten = flatten

        self.proj = nn.Conv2d(
            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
        )
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

        self.position_getter = PositionGetter()

    def forward(self, x):
        B, C, H, W = x.shape
        torch._assert(
            H == self.img_size[0],
            f"Input image height ({H}) doesn't match model ({self.img_size[0]}).",
        )
        torch._assert(
            W == self.img_size[1],
            f"Input image width ({W}) doesn't match model ({self.img_size[1]}).",
        )
        x = self.proj(x)
        pos = self.position_getter(B, x.size(2), x.size(3), x.device)
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
        x = self.norm(x)
        return x, pos

    def _init_weights(self):
        w = self.proj.weight.data


forward

        embed_dim=768,
        norm_layer=None,
        flatten=True,
    ):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
        self.num_patches = self.grid_size[0] * self.grid_size[1]
        self.flatten = flatten

        self.proj = nn.Conv2d(
            in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
        )
        self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()

        self.position_getter = PositionGetter()

    def forward(self, x):
        B, C, H, W = x.shape
        torch._assert(
            H == self.img_size[0],
            f"Input image height ({H}) doesn't match model ({self.img_size[0]}).",
        )
        torch._assert(
            W == self.img_size[1],
            f"Input image width ({W}) doesn't match model ({self.img_size[1]}).",
        )
        x = self.proj(x)
        pos = self.position_getter(B, x.size(2), x.size(3), x.device)
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
        x = self.norm(x)
        return x, pos

    def _init_weights(self):
        w = self.proj.weight.data
        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))

if __name__ == "__main__":
    import os
    import sys

    sys.path.append(os.path.dirname(os.path.dirname(__file__)))
    import dust3r.utils.path_to_croco
    from models.pos_embed import get_2d_sincos_pos_embed, RoPE2D
    from functools import partial
    from torch.utils.checkpoint import checkpoint

    torch.manual_seed(0)

    enc_blocks_ray_map = (
        nn.ModuleList(


    shape = (x.shape[0],) + (1,) * (
        x.ndim - 1
    )  # work with diff dim tensors, not just 2D ConvNets
    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
    if keep_prob > 0.0 and scale_by_keep:
        random_tensor.div_(keep_prob)
    return x * random_tensor

class DropPath(nn.Module):
    """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks)."""

    def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob
        self.scale_by_keep = scale_by_keep

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)

    def extra_repr(self):
        return f"drop_prob={round(self.drop_prob,3):0.3f}"

class Mlp(nn.Module):
    """MLP as used in Vision Transformer, MLP-Mixer and related networks"""

    def __init__(
        self,
        in_features,
        hidden_features=None,
        out_features=None,
        act_layer=nn.GELU,
        bias=True,
        drop=0.0,
    ):
        super().__init__()
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features
        bias = to_2tuple(bias)
        drop_probs = to_2tuple(drop)


modulate

        return x, y

class ModLN(nn.Module):
    """
    Modulation with adaLN.

    References:
    DiT: https://github.com/facebookresearch/DiT/blob/main/models.py#L101
    """

    def __init__(self, inner_dim: int, mod_dim: int, eps: float):
        super().__init__()
        self.norm = nn.LayerNorm(inner_dim, eps=eps)
        self.mlp = nn.Sequential(
            nn.SiLU(),
            nn.Linear(mod_dim, inner_dim * 2),
        )

    @staticmethod
    def modulate(x, shift, scale):

        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)

    def forward(self, x: torch.Tensor, mod: torch.Tensor) -> torch.Tensor:
        shift, scale = self.mlp(mod).chunk(2, dim=-1)  # [N, D]
        return self.modulate(self.norm(x), shift, scale)  # [N, L, D]

class ConditionModulationBlock(nn.Module):

    def __init__(
        self,
        dim,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=False,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=partial(ModLN, eps=1e-6),
        rope=None,


        mlp_hidden_dim = int(dim * mlp_ratio)
        self.mlp = Mlp(
            in_features=dim,
            hidden_features=mlp_hidden_dim,
            act_layer=act_layer,
            drop=drop,
        )

    def forward(self, x, mod, xpos):
        x = x + self.drop_path(self.attn(self.norm1(x, mod), xpos))
        x = x + self.drop_path(self.mlp(self.norm2(x, mod)))
        return x

class PositionGetter(object):
    """return positions of patches"""

    def __init__(self):
        self.cache_positions = {}

    def __call__(self, b, h, w, device):
        if not (h, w) in self.cache_positions:
            x = torch.arange(w, device=device)
            y = torch.arange(h, device=device)
            self.cache_positions[h, w] = torch.cartesian_prod(y, x)  # (h, w, 2)
        pos = self.cache_positions[h, w].view(1, h * w, 2).expand(b, -1, 2).clone()
        return pos

class PatchEmbed(nn.Module):
    """just adding _init_weights + position getter compared to timm.models.layers.patch_embed.PatchEmbed"""

    def __init__(
        self,
        img_size=224,
        patch_size=16,
        in_chans=3,
        embed_dim=768,
        norm_layer=None,
        flatten=True,
    ):
        super().__init__()
        img_size = to_2tuple(img_size)
        patch_size = to_2tuple(patch_size)
        self.img_size = img_size
        self.patch_size = patch_size
        self.grid_size = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])


        self.position_getter = PositionGetter()

    def forward(self, x):
        B, C, H, W = x.shape
        torch._assert(
            H == self.img_size[0],
            f"Input image height ({H}) doesn't match model ({self.img_size[0]}).",
        )
        torch._assert(
            W == self.img_size[1],
            f"Input image width ({W}) doesn't match model ({self.img_size[1]}).",
        )
        x = self.proj(x)
        pos = self.position_getter(B, x.size(2), x.size(3), x.device)
        if self.flatten:
            x = x.flatten(2).transpose(1, 2)  # BCHW -> BNC
        x = self.norm(x)
        return x, pos

    def _init_weights(self):
        w = self.proj.weight.data
        torch.nn.init.xavier_uniform_(w.view([w.shape[0], -1]))

if __name__ == "__main__":
    import os
    import sys

    sys.path.append(os.path.dirname(os.path.dirname(__file__)))
    import dust3r.utils.path_to_croco
    from models.pos_embed import get_2d_sincos_pos_embed, RoPE2D
    from functools import partial
    from torch.utils.checkpoint import checkpoint

    torch.manual_seed(0)

    enc_blocks_ray_map = (
        nn.ModuleList(
            [
                Block(
                    768,
                    16,


from copy import copy, deepcopy
import torch
import torch.nn as nn
import torch.nn.functional as F

from dust3r.utils.geometry import (
    inv,
    geotrf,
    normalize_pointcloud_group,
    get_group_pointcloud_center_scale,
    to_euclidean_dist,
)
import numpy as np
from dust3r.utils.camera import (
    pose_encoding_to_camera,
    camera_to_pose_encoding,
    relative_pose_absT_quatR,
)
from dust3r.utils.image import unpad_image
from dust3r.utils import SMPL_Layer

import roma
from tqdm import tqdm

def Sum(*losses_and_masks):
    loss, mask = losses_and_masks[0]
    if loss.ndim > 0:
        # we are actually returning the loss for every pixels
        return losses_and_masks
    else:
        # we are returning the global loss
        for loss2, mask2 in losses_and_masks[1:]:
            loss = loss + loss2
        return loss

def stack_view(ls, k=None):
    if isinstance(ls[0], dict):
        v = torch.stack([g[k] for g in ls], dim=0)
    else:
        v = torch.stack(ls, dim=0)
    return v.view(-1, *v.shape[2:])

def _neg_loss(pred, gt):
  '''
  Code modified from: https://github.com/xingyizhou/CenterNet/blob/4c50fd3a46bdf63dbf2082c5cbb3458d39579e6c/src/lib/models/losses.py#L42
    Modified focal loss. Exactly the same as CornerNet.
      Runs faster and costs a little bit more memory
    Arguments:
      pred (batch x c x h x w)
      gt_regr (batch x c x h x w)

  '''
  assert pred.shape == gt.shape

  pos_inds = gt.eq(1).float()
  neg_inds = gt.lt(1).float()

  neg_weights = torch.pow(1 - gt, 4)

  loss = 0

  eps = 1e-7

  pos_loss = torch.log(pred + eps) * torch.pow(1 - pred, 2) * pos_inds
  neg_loss = torch.log(1 - pred + eps) * torch.pow(pred, 2) * neg_weights * neg_inds

  num_pos  = pos_inds.float().sum()
  pos_loss = pos_loss.sum()
  neg_loss = neg_loss.sum()

  if num_pos == 0:
    loss = loss - neg_loss
  else:
    loss = loss - (pos_loss + neg_loss) / num_pos
  return loss

class BaseCriterion(nn.Module):
    def __init__(self, reduction="mean"):
        super().__init__()
        self.reduction = reduction

class LLoss(BaseCriterion):
    """L-norm loss"""

    def forward(self, a, b):
        # assert (
        #     a.shape == b.shape and a.ndim >= 2 and 1 <= a.shape[-1] <= 3
        # ), f"Bad shape = {a.shape}"
        dist = self.distance(a, b)
        if self.reduction == "none":
            return dist
        if self.reduction == "sum":
            return dist.sum()
        if self.reduction == "mean":
            return dist.mean() if dist.numel() > 0 else dist.new_zeros(())
        raise ValueError(f"bad {self.reduction=} mode")

    def distance(self, a, b):
        raise NotImplementedError()

class L21Loss(LLoss):
    """Euclidean distance between 3d points"""

    def distance(self, a, b):
        return torch.norm(a - b, dim=-1)  # normalized L2 distance

L21 = L21Loss()

class MSELoss(LLoss):
    def distance(self, a, b):
        return (a - b) ** 2

MSE = MSELoss()

class L1Loss(LLoss):
    def distance(self, a, b):
        return (a - b).abs()

L1 = L1Loss()

class Criterion(nn.Module):
    def __init__(self, criterion=None):
        super().__init__()
        assert isinstance(
            criterion, BaseCriterion
        ), f"{criterion} is not a proper criterion!"
        self.criterion = copy(criterion)

    def get_name(self):
        return f"{type(self).__name__}({self.criterion})"

    def with_reduction(self, mode="none"):
        res = loss = deepcopy(self)
        while loss is not None:
            assert isinstance(loss, Criterion)
            loss.criterion.reduction = mode  # make it return the loss for each sample
            loss = loss._loss2  # we assume loss is a Multiloss
        return res

class MultiLoss(nn.Module):
    """Easily combinable losses (also keep track of individual loss values):
        loss = MyLoss1() + 0.1*MyLoss2()
    Usage:
        Inherit from this class and override get_name() and compute_loss()
    """

    def __init__(self):
        super().__init__()
        self._alpha = 1
        self._loss2 = None

    def compute_loss(self, *args, **kwargs):
        raise NotImplementedError()

    def get_name(self):
        raise NotImplementedError()

    def __mul__(self, alpha):
        assert isinstance(alpha, (int, float))
        res = copy(self)
        res._alpha = alpha
        return res

    __rmul__ = __mul__  # same

    def __add__(self, loss2):
        assert isinstance(loss2, MultiLoss)
        res = cur = copy(self)
        # find the end of the chain
        while cur._loss2 is not None:
            cur = cur._loss2
        cur._loss2 = loss2
        return res

    def __repr__(self):
        name = self.get_name()
        if self._alpha != 1:
            name = f"{self._alpha:g}*{name}"
        if self._loss2:
            name = f"{name} + {self._loss2}"
        return name

    def forward(self, *args, **kwargs):
        loss = self.compute_loss(*args, **kwargs)
        if isinstance(loss, tuple):
            loss, details = loss
        elif loss.ndim == 0:
            details = {self.get_name(): float(loss)}
        else:
            details = {}
        loss = loss * self._alpha

        if self._loss2:
            loss2, details2 = self._loss2(*args, **kwargs)
            loss = loss + loss2
            details |= details2

        return loss, details

class SSIM(nn.Module):
    """Layer to compute the SSIM loss between a pair of images"""

    def __init__(self):
        super(SSIM, self).__init__()
        self.mu_x_pool = nn.AvgPool2d(3, 1)
        self.mu_y_pool = nn.AvgPool2d(3, 1)
        self.sig_x_pool = nn.AvgPool2d(3, 1)
        self.sig_y_pool = nn.AvgPool2d(3, 1)
        self.sig_xy_pool = nn.AvgPool2d(3, 1)

        self.refl = nn.ReflectionPad2d(1)

        self.C1 = 0.01**2
        self.C2 = 0.03**2

    def forward(self, x, y):
        x = self.refl(x)
        y = self.refl(y)

        mu_x = self.mu_x_pool(x)
        mu_y = self.mu_y_pool(y)

        sigma_x = self.sig_x_pool(x**2) - mu_x**2
        sigma_y = self.sig_y_pool(y**2) - mu_y**2
        sigma_xy = self.sig_xy_pool(x * y) - mu_x * mu_y

        SSIM_n = (2 * mu_x * mu_y + self.C1) * (2 * sigma_xy + self.C2)
        SSIM_d = (mu_x**2 + mu_y**2 + self.C1) * (sigma_x + sigma_y + self.C2)

        return torch.clamp((1 - SSIM_n / SSIM_d) / 2, 0, 1)

class RGBLoss(Criterion, MultiLoss):
    def __init__(self, criterion):
        super().__init__(criterion)
        self.ssim = SSIM()

    def img_loss(self, a, b):
        return self.criterion(a, b)

    def compute_loss(self, gts, preds, **kw):
        gt_rgbs = [gt["img"].permute(0, 2, 3, 1) for gt in gts]
        pred_rgbs = [pred["rgb"] for pred in preds]
        ls = [
            self.img_loss(pred_rgb, gt_rgb)
            for pred_rgb, gt_rgb in zip(pred_rgbs, gt_rgbs)
        ]
        details = {}
        self_name = type(self).__name__
        for i, l in enumerate(ls):
            details[self_name + f"_rgb/{i+1}"] = float(l)
            details[f"pred_rgb_{i+1}"] = pred_rgbs[i]
        rgb_loss = sum(ls) / len(ls)
        return rgb_loss, details

class SMPLLoss(Criterion, MultiLoss):
    def __init__(self, criterion):
        super().__init__(criterion)
        scale = 0.01
        self.alpha_msk = 100.0 * scale
        self.alpha_bce = 10.0 * scale
        self.alpha_rotmat = 100.0 * scale
        self.alpha_shape = 10.0 * scale
        self.alpha_transl = 100.0 * scale
        self.alpha_j3d = 100.0 * scale
        self.alpha_v3d = 100.0 * scale
        self.alpha_j2d = 1.0 * scale
        self.alpha_v2d = 1.0 * scale

        # SMPL layer
        person_center = 'head'
        dict_smpl_layer = {
            'neutral': {
                10: SMPL_Layer(type='smplx', gender='neutral', num_betas=10, kid=False, person_center=person_center),
                11: SMPL_Layer(type='smplx', gender='neutral', num_betas=11, kid=False, person_center=person_center),
                }
            }
        _moduleDict = []
        for k, _smpl_layer in dict_smpl_layer.items():
            for x, y in _smpl_layer.items():
                _moduleDict.append([f"{k}_{x}", deepcopy(y)])
        self.smpl_layer = nn.ModuleDict(_moduleDict)

    def get_name(self):
        return "SMPLLoss"

    def mask_loss(self, gts, preds, masks, ret_pred=False):
        gt_msks = [gt["msk_mhmr"].unsqueeze(-1) for gt in gts]
        pred_msks = [pred["msk"] for pred in preds]
        ls = [
            F.binary_cross_entropy(p[m], g[m])
            for p, g, m in zip(pred_msks, gt_msks, masks)
        ]
        details = {}
        self_name = self.get_name()
        for i, l in enumerate(ls):
            details[self_name + f"_msk/{i+1}"] = float(l)
            if ret_pred:
                details[f"pred_msk_{i+1}"] = pred_msks[i]
        bce = sum(ls) / len(ls)
        return bce, details

    def bce(self, gts, preds, masks, ret_pred=False):
        gt_scores = [(gt["smpl_scores"] >= 1).to(int).unsqueeze(-1) for gt in gts]
        pred_scores = [pred["smpl_scores"] for pred in preds]
        ls = [
            _neg_loss(p[m], g[m])
            for p, g, m in zip(pred_scores, gt_scores, masks)
        ]
        details = {}
        self_name = self.get_name()
        for i, l in enumerate(ls):
            details[self_name + f"_scores/{i+1}"] = float(l)
            if ret_pred:
                details[f"pred_smpl_scores_{i+1}"] = pred_scores[i]
        bce = sum(ls) / len(ls)
        return bce, details

    def smpl_param_loss(self, gts, preds, masks, k, ret_pred=False):
        if isinstance(gts[0], dict):
            gts = [gt[k] for gt in gts]
        if isinstance(preds[0], dict):
            preds = [pred[k] for pred in preds]

        ls = [
            self.criterion(p[m], g[m])
            for p, g, m in zip(preds, gts, masks)
        ]
        details = {}
        self_name = self.get_name()
        k_name = k.split('smpl_')[-1] if 'smpl_' in k else k
        for i, l in enumerate(ls):
            details[self_name + f"_{k_name}/{i+1}"] = float(l)
            if ret_pred:
                details[f"pred_{k}_{i+1}"] = preds[i]
        loss = sum(ls) / len(ls)

        return loss, details

    def point3d_loss(self, gts, preds, gt_t_p, pr_t_ps, masks, k, ret_pred=False):
        if isinstance(gts[0], dict):
            gts = [gt[k] for gt in gts]
        if isinstance(preds[0], dict):
            preds = [pred[k] for pred in preds]

        if gt_t_p is None or pr_t_ps is None:
            ls = [
                self.criterion(p[m], g[m])
                for p, g, m in zip(preds, gts, masks)
            ]
        else:
            ls = [
                self.criterion(p[m]-pr_p[m], g[m]-gt_p[m])
                for p, pr_p, g, gt_p, m in zip(preds, pr_t_ps, gts, gt_t_p, masks)
            ]

        details = {}
        self_name = self.get_name()
        k_name = k.split('smpl_')[-1] if 'smpl_' in k else k
        if gt_t_p is None or pr_t_ps is None:
            k_name = "c" + k_name
        for i, l in enumerate(ls):
            details[self_name + f"_{k_name}/{i+1}"] = float(l)
            if ret_pred:
                details[f"pred_{k}_{i+1}"] = preds[i]
        loss = sum(ls) / len(ls)

        return loss, details

    def point2d_loss(self, gts, preds, masks, k, shape=None, ret_pred=False):
        if isinstance(gts[0], dict):
            shape = gts[0]['true_shape'][0]
            gts = [gt[k] for gt in gts]
        if isinstance(preds[0], dict):
            preds = [pred[k] for pred in preds]

        valid_mask = [
            ((gt[..., 0] > 0) & (gt[..., 0] < shape[1]) & (gt[..., 1] > 0) & (gt[..., 1] < shape[0]
            )) for gt in gts]

        ls = [
            self.criterion(p[m1.unsqueeze(-1) & m2], g[m1.unsqueeze(-1) & m2])
            for p, g, m1, m2 in zip(preds, gts, masks, valid_mask)
        ]
        details = {}
        self_name = self.get_name()
        k_name = k.split('smpl_')[-1] if 'smpl_' in k else k
        for i, l in enumerate(ls):
            details[self_name + f"_{k_name}/{i+1}"] = float(l)
            if ret_pred:
                details[f"pred_{k}_{i+1}"] = preds[i]
        loss = sum(ls) / len(ls)

        return loss, details

    def compute_loss(self, gts, preds, **kw):
        img_mask_list = [gt["img_mask"] for gt in gts]
        smpl_mask_list = [gt["smpl_mask"] for gt in gts]
        masks_list = [a.unsqueeze(1) & b for a, b in zip(img_mask_list, smpl_mask_list)]

        # Detection loss
        score_loss, score_details = self.bce(gts, preds, img_mask_list, ret_pred=True)

        has_msk = "msk" in preds[0]
        if has_msk:
            msk_loss, msk_details = self.mask_loss(gts, preds, img_mask_list, ret_pred=True)

        K = stack_view(gts, 'camera_intrinsics')
        img_mask = stack_view(img_mask_list,'img_mask').unsqueeze(1)
        smpl_mask = stack_view(smpl_mask_list, 'smpl_mask') * img_mask
        idx_h = torch.where(smpl_mask)
        if int(smpl_mask.sum()) == 0:
            total_loss = self.alpha_bce * score_loss
            details = {
                **score_details,
            }
            if has_msk:
                total_loss += self.alpha_msk * msk_loss
                details.update(msk_details)
            return total_loss, details

        # Prediction
        pred_rotmat = stack_view(preds, 'smpl_rotmat')
        pred_rotvec = roma.rotmat_to_rotvec(pred_rotmat[smpl_mask])
        pred_shape = stack_view(preds, 'smpl_shape')
        pred_transl = [pred.pop("smpl_transl") for pred in preds]
        pred_transl = stack_view(pred_transl, 'smpl_transl')
        pred_expression = stack_view(preds, 'smpl_expression')

        smpl_out = self.smpl_layer[f"neutral_{pred_shape.shape[-1]}"](
            pred_rotvec,
            pred_shape[smpl_mask],
            pred_transl[smpl_mask],
            None, None,
            K=K[idx_h[0]],
            expression=pred_expression[smpl_mask])

        pred_smpl = {}
        batch_size = img_mask_list[0].shape[0]
        num_view = len(gts)
        max_humans = smpl_mask.shape[1]

        for k, v in smpl_out.items():
            full_out = torch.zeros(
                num_view * batch_size, max_humans, *v.shape[1:],
                device=v.device, dtype=v.dtype,
            )
            full_out[smpl_mask] = v
            pred_smpl[k] = full_out.chunk(num_view, dim=0)

        # SMPL-X params
        rotmat_loss, rotmat_details = self.smpl_param_loss(gts, preds, masks_list, "smpl_rotmat")
        transl_loss, transl_details = self.smpl_param_loss(gts, pred_smpl['smpl_transl'], masks_list, "smpl_transl")
        shape_dim = min([gts[0]['smpl_shape'].shape[-1], preds[0]['smpl_shape'].shape[-1]])
        gt_shape = [gt['smpl_shape'][...,:shape_dim] for gt in gts]
        pred_shape = [pred['smpl_shape'][...,:shape_dim] for pred in preds]
        shape_loss, shape_details = self.smpl_param_loss(gt_shape, pred_shape, masks_list, "smpl_shape")

        # 3D points
        gt_transl_pelvis= [gt['smpl_transl_pelvis'][..., None, :] for gt in gts]
        pred_transl_pelvis = pred_smpl['smpl_transl_pelvis']
        j3d_loss, j3d_details = self.point3d_loss(
                                    gts,
                                    pred_smpl['smpl_j3d'],
                                    gt_transl_pelvis,
                                    pred_transl_pelvis,
                                    masks_list, "smpl_j3d")
        v3d_loss, v3d_details = self.point3d_loss(
                                    gts,
                                    pred_smpl['smpl_v3d'],
                                    gt_transl_pelvis,
                                    pred_transl_pelvis,
                                    masks_list, "smpl_v3d",
                                    ret_pred=True)
        cj3d_loss, cj3d_details = self.point3d_loss(
                                    gts,
                                    pred_smpl['smpl_j3d'],
                                    None,
                                    None,
                                    masks_list, "smpl_j3d")
        cv3d_loss, cv3d_details = self.point3d_loss(
                                    gts,
                                    pred_smpl['smpl_v3d'],
                                    None,
                                    None,
                                    masks_list, "smpl_v3d")

        # total loss
        total_loss = self.alpha_bce * score_loss +\
                     self.alpha_rotmat * rotmat_loss +\
                     self.alpha_shape * shape_loss +\
                     self.alpha_transl * transl_loss +\
                     self.alpha_j3d * j3d_loss +\
                     self.alpha_v3d * v3d_loss +\
                     self.alpha_j3d * cj3d_loss +\
                     self.alpha_v3d * cv3d_loss

        details = {
            **score_details,
            **rotmat_details,
            **transl_details,
            **shape_details,
            **j3d_details,
            **v3d_details,
            **cj3d_details,
            **cv3d_details,
        }

        if has_msk:
            total_loss += self.alpha_msk * msk_loss
            details.update(msk_details)

        if cv3d_loss < 1.0:
            # 2D reprojection
            j2d_loss, j2d_details = self.point2d_loss(
                                        gts,
                                        pred_smpl['smpl_j2d'],
                                        masks_list, "smpl_j2d")
            v2d_loss, v2d_details = self.point2d_loss(
                                        gts,
                                        pred_smpl['smpl_v2d'],
                                        masks_list, "smpl_v2d")
            total_loss += self.alpha_j2d * j2d_loss +\
                          self.alpha_v2d * v2d_loss
            details.update({
                **j2d_details,
                **v2d_details
            })

        return total_loss, details

class NaiveSMPLLoss(SMPLLoss):
    def __init__(self, criterion):
        super().__init__(criterion)

    def compute_loss(self, gts, preds, **kw):
        img_mask_list = [gt["img_mask"] for gt in gts]
        smpl_mask_list = [gt["smpl_mask"] for gt in gts]
        masks_list = [a.unsqueeze(1) & b for a, b in zip(img_mask_list, smpl_mask_list)]

        # Detection loss
        score_loss, score_details = self.bce(gts, preds, img_mask_list, ret_pred=True)

        # only for inference SMPL model
        K = stack_view(gts, 'camera_intrinsics')
        img_mask = stack_view(img_mask_list,'img_mask').unsqueeze(1)
        smpl_mask = stack_view(smpl_mask_list, 'smpl_mask') * img_mask
        idx_h = torch.where(smpl_mask)
        if int(smpl_mask.sum()) == 0:
            total_loss = self.alpha_bce * score_loss
            details = {
                **score_details,
            }
            return total_loss, details

        # Prediction
        pred_rotmat = stack_view(preds, 'smpl_rotmat')
        pred_rotvec = roma.rotmat_to_rotvec(pred_rotmat[smpl_mask])
        pred_shape = stack_view(preds, 'smpl_shape')
        pred_transl = [pred.pop("smpl_transl") for pred in preds]
        pred_transl = stack_view(pred_transl, 'smpl_transl')
        pred_expression = stack_view(preds, 'smpl_expression')

        # Neutral for MHMR
        K_mhmr = stack_view(gts, 'K_mhmr')
        mhmr_img_res = gts[0]["img_mhmr"].shape[-1]
        # fine head uv
        pred_loc = stack_view(preds, 'smpl_loc')
        # Distance
        dist = pred_transl[smpl_mask][:, 0].unsqueeze(-1)
        dist = to_euclidean_dist(mhmr_img_res, dist, K_mhmr[idx_h[0]])  # use K GT
        smpl_out = self.smpl_layer[f"neutral_{pred_shape.shape[-1]}"](
            pred_rotvec,
            pred_shape[smpl_mask],
            None,
            pred_loc[smpl_mask],
            dist,
# ... truncated ...


Sum

from dust3r.utils.geometry import (
    inv,
    geotrf,
    normalize_pointcloud_group,
    get_group_pointcloud_center_scale,
    to_euclidean_dist,
)
import numpy as np
from dust3r.utils.camera import (
    pose_encoding_to_camera,
    camera_to_pose_encoding,
    relative_pose_absT_quatR,
)
from dust3r.utils.image import unpad_image
from dust3r.utils import SMPL_Layer

import roma
from tqdm import tqdm

def Sum(*losses_and_masks):
    loss, mask = losses_and_masks[0]
    if loss.ndim > 0:
        # we are actually returning the loss for every pixels
        return losses_and_masks
    else:
        # we are returning the global loss
        for loss2, mask2 in losses_and_masks[1:]:
            loss = loss + loss2
        return loss

def stack_view(ls, k=None):
    if isinstance(ls[0], dict):
        v = torch.stack([g[k] for g in ls], dim=0)
    else:
        v = torch.stack(ls, dim=0)
    return v.view(-1, *v.shape[2:])

def _neg_loss(pred, gt):
  '''
  Code modified from: https://github.com/xingyizhou/CenterNet/blob/4c50fd3a46bdf63dbf2082c5cbb3458d39579e6c/src/lib/models/losses.py#L42
    Modified focal loss. Exactly the same as CornerNet.
      Runs faster and costs a little bit more memory
    Arguments:
      pred (batch x c x h x w)
      gt_regr (batch x c x h x w)

  '''


    relative_pose_absT_quatR,
)
from dust3r.utils.image import unpad_image
from dust3r.utils import SMPL_Layer

import roma
from tqdm import tqdm

def Sum(*losses_and_masks):
    loss, mask = losses_and_masks[0]
    if loss.ndim > 0:
        # we are actually returning the loss for every pixels
        return losses_and_masks
    else:
        # we are returning the global loss
        for loss2, mask2 in losses_and_masks[1:]:
            loss = loss + loss2
        return loss

def stack_view(ls, k=None):
    if isinstance(ls[0], dict):
        v = torch.stack([g[k] for g in ls], dim=0)
    else:
        v = torch.stack(ls, dim=0)
    return v.view(-1, *v.shape[2:])

def _neg_loss(pred, gt):
  '''
  Code modified from: https://github.com/xingyizhou/CenterNet/blob/4c50fd3a46bdf63dbf2082c5cbb3458d39579e6c/src/lib/models/losses.py#L42
    Modified focal loss. Exactly the same as CornerNet.
      Runs faster and costs a little bit more memory
    Arguments:
      pred (batch x c x h x w)
      gt_regr (batch x c x h x w)

  '''
  assert pred.shape == gt.shape

  pos_inds = gt.eq(1).float()
  neg_inds = gt.lt(1).float()

  neg_weights = torch.pow(1 - gt, 4)

  loss = 0


def Sum(*losses_and_masks):
    loss, mask = losses_and_masks[0]
    if loss.ndim > 0:
        # we are actually returning the loss for every pixels
        return losses_and_masks
    else:
        # we are returning the global loss
        for loss2, mask2 in losses_and_masks[1:]:
            loss = loss + loss2
        return loss

def stack_view(ls, k=None):
    if isinstance(ls[0], dict):
        v = torch.stack([g[k] for g in ls], dim=0)
    else:
        v = torch.stack(ls, dim=0)
    return v.view(-1, *v.shape[2:])

def _neg_loss(pred, gt):
  '''
  Code modified from: https://github.com/xingyizhou/CenterNet/blob/4c50fd3a46bdf63dbf2082c5cbb3458d39579e6c/src/lib/models/losses.py#L42
    Modified focal loss. Exactly the same as CornerNet.
      Runs faster and costs a little bit more memory
    Arguments:
      pred (batch x c x h x w)
      gt_regr (batch x c x h x w)

  '''
  assert pred.shape == gt.shape

  pos_inds = gt.eq(1).float()
  neg_inds = gt.lt(1).float()

  neg_weights = torch.pow(1 - gt, 4)

  loss = 0

  eps = 1e-7

  pos_loss = torch.log(pred + eps) * torch.pow(1 - pred, 2) * pos_inds
  neg_loss = torch.log(1 - pred + eps) * torch.pow(pred, 2) * neg_weights * neg_inds

  num_pos  = pos_inds.float().sum()
  pos_loss = pos_loss.sum()
  neg_loss = neg_loss.sum()

  if num_pos == 0:
    loss = loss - neg_loss
  else:
    loss = loss - (pos_loss + neg_loss) / num_pos
  return loss

class BaseCriterion(nn.Module):
    def __init__(self, reduction="mean"):
        super().__init__()
        self.reduction = reduction

class LLoss(BaseCriterion):
    """L-norm loss"""

    def forward(self, a, b):
        # assert (
        #     a.shape == b.shape and a.ndim >= 2 and 1 <= a.shape[-1] <= 3
        # ), f"Bad shape = {a.shape}"
        dist = self.distance(a, b)
        if self.reduction == "none":
            return dist
        if self.reduction == "sum":
            return dist.sum()


BaseCriterion

  neg_weights = torch.pow(1 - gt, 4)

  loss = 0

  eps = 1e-7

  pos_loss = torch.log(pred + eps) * torch.pow(1 - pred, 2) * pos_inds
  neg_loss = torch.log(1 - pred + eps) * torch.pow(pred, 2) * neg_weights * neg_inds

  num_pos  = pos_inds.float().sum()
  pos_loss = pos_loss.sum()
  neg_loss = neg_loss.sum()

  if num_pos == 0:
    loss = loss - neg_loss
  else:
    loss = loss - (pos_loss + neg_loss) / num_pos
  return loss

class BaseCriterion(nn.Module):
    def __init__(self, reduction="mean"):
        super().__init__()
        self.reduction = reduction

class LLoss(BaseCriterion):
    """L-norm loss"""

    def forward(self, a, b):
        # assert (
        #     a.shape == b.shape and a.ndim >= 2 and 1 <= a.shape[-1] <= 3
        # ), f"Bad shape = {a.shape}"
        dist = self.distance(a, b)
        if self.reduction == "none":
            return dist
        if self.reduction == "sum":
            return dist.sum()
        if self.reduction == "mean":
            return dist.mean() if dist.numel() > 0 else dist.new_zeros(())
        raise ValueError(f"bad {self.reduction=} mode")

    def distance(self, a, b):
        raise NotImplementedError()


LLoss

  pos_loss = torch.log(pred + eps) * torch.pow(1 - pred, 2) * pos_inds
  neg_loss = torch.log(1 - pred + eps) * torch.pow(pred, 2) * neg_weights * neg_inds

  num_pos  = pos_inds.float().sum()
  pos_loss = pos_loss.sum()
  neg_loss = neg_loss.sum()

  if num_pos == 0:
    loss = loss - neg_loss
  else:
    loss = loss - (pos_loss + neg_loss) / num_pos
  return loss

class BaseCriterion(nn.Module):
    def __init__(self, reduction="mean"):
        super().__init__()
        self.reduction = reduction

class LLoss(BaseCriterion):
    """L-norm loss"""

    def forward(self, a, b):
        # assert (
        #     a.shape == b.shape and a.ndim >= 2 and 1 <= a.shape[-1] <= 3
        # ), f"Bad shape = {a.shape}"
        dist = self.distance(a, b)
        if self.reduction == "none":
            return dist
        if self.reduction == "sum":
            return dist.sum()
        if self.reduction == "mean":
            return dist.mean() if dist.numel() > 0 else dist.new_zeros(())
        raise ValueError(f"bad {self.reduction=} mode")

    def distance(self, a, b):
        raise NotImplementedError()

class L21Loss(LLoss):
    """Euclidean distance between 3d points"""

    def distance(self, a, b):
        return torch.norm(a - b, dim=-1)  # normalized L2 distance

L21 = L21Loss()

class MSELoss(LLoss):
    def distance(self, a, b):
        return (a - b) ** 2

MSE = MSELoss()


class LLoss(BaseCriterion):
    """L-norm loss"""

    def forward(self, a, b):
        # assert (
        #     a.shape == b.shape and a.ndim >= 2 and 1 <= a.shape[-1] <= 3
        # ), f"Bad shape = {a.shape}"
        dist = self.distance(a, b)
        if self.reduction == "none":
            return dist
        if self.reduction == "sum":
            return dist.sum()
        if self.reduction == "mean":
            return dist.mean() if dist.numel() > 0 else dist.new_zeros(())
        raise ValueError(f"bad {self.reduction=} mode")

    def distance(self, a, b):
        raise NotImplementedError()

class L21Loss(LLoss):
    """Euclidean distance between 3d points"""

    def distance(self, a, b):
        return torch.norm(a - b, dim=-1)  # normalized L2 distance

L21 = L21Loss()

class MSELoss(LLoss):
    def distance(self, a, b):
        return (a - b) ** 2

MSE = MSELoss()

class L1Loss(LLoss):
    def distance(self, a, b):
        return (a - b).abs()

L1 = L1Loss()


MSELoss

        if self.reduction == "sum":
            return dist.sum()
        if self.reduction == "mean":
            return dist.mean() if dist.numel() > 0 else dist.new_zeros(())
        raise ValueError(f"bad {self.reduction=} mode")

    def distance(self, a, b):
        raise NotImplementedError()

class L21Loss(LLoss):
    """Euclidean distance between 3d points"""

    def distance(self, a, b):
        return torch.norm(a - b, dim=-1)  # normalized L2 distance

L21 = L21Loss()

class MSELoss(LLoss):
    def distance(self, a, b):
        return (a - b) ** 2

MSE = MSELoss()

class L1Loss(LLoss):
    def distance(self, a, b):
        return (a - b).abs()

L1 = L1Loss()

class Criterion(nn.Module):
    def __init__(self, criterion=None):
        super().__init__()
        assert isinstance(
            criterion, BaseCriterion
        ), f"{criterion} is not a proper criterion!"
        self.criterion = copy(criterion)


class L21Loss(LLoss):
    """Euclidean distance between 3d points"""

    def distance(self, a, b):
        return torch.norm(a - b, dim=-1)  # normalized L2 distance

L21 = L21Loss()

class MSELoss(LLoss):
    def distance(self, a, b):
        return (a - b) ** 2

MSE = MSELoss()

class L1Loss(LLoss):
    def distance(self, a, b):
        return (a - b).abs()

L1 = L1Loss()

class Criterion(nn.Module):
    def __init__(self, criterion=None):
        super().__init__()
        assert isinstance(
            criterion, BaseCriterion
        ), f"{criterion} is not a proper criterion!"
        self.criterion = copy(criterion)

    def get_name(self):
        return f"{type(self).__name__}({self.criterion})"

    def with_reduction(self, mode="none"):
        res = loss = deepcopy(self)
        while loss is not None:
            assert isinstance(loss, Criterion)
            loss.criterion.reduction = mode  # make it return the loss for each sample


Criterion

L21 = L21Loss()

class MSELoss(LLoss):
    def distance(self, a, b):
        return (a - b) ** 2

MSE = MSELoss()

class L1Loss(LLoss):
    def distance(self, a, b):
        return (a - b).abs()

L1 = L1Loss()

class Criterion(nn.Module):
    def __init__(self, criterion=None):
        super().__init__()
        assert isinstance(
            criterion, BaseCriterion
        ), f"{criterion} is not a proper criterion!"
        self.criterion = copy(criterion)

    def get_name(self):
        return f"{type(self).__name__}({self.criterion})"

    def with_reduction(self, mode="none"):
        res = loss = deepcopy(self)
        while loss is not None:
            assert isinstance(loss, Criterion)
            loss.criterion.reduction = mode  # make it return the loss for each sample
            loss = loss._loss2  # we assume loss is a Multiloss
        return res

class MultiLoss(nn.Module):
    """Easily combinable losses (also keep track of individual loss values):
        loss = MyLoss1() + 0.1*MyLoss2()
    Usage:
        Inherit from this class and override get_name() and compute_loss()
    """

    def __init__(self):
        super().__init__()
        self._alpha = 1
        self._loss2 = None

    def compute_loss(self, *args, **kwargs):
        raise NotImplementedError()

    def get_name(self):
        raise NotImplementedError()


MultiLoss

class Criterion(nn.Module):
    def __init__(self, criterion=None):
        super().__init__()
        assert isinstance(
            criterion, BaseCriterion
        ), f"{criterion} is not a proper criterion!"
        self.criterion = copy(criterion)

    def get_name(self):
        return f"{type(self).__name__}({self.criterion})"

    def with_reduction(self, mode="none"):
        res = loss = deepcopy(self)
        while loss is not None:
            assert isinstance(loss, Criterion)
            loss.criterion.reduction = mode  # make it return the loss for each sample
            loss = loss._loss2  # we assume loss is a Multiloss
        return res

class MultiLoss(nn.Module):
    """Easily combinable losses (also keep track of individual loss values):
        loss = MyLoss1() + 0.1*MyLoss2()
    Usage:
        Inherit from this class and override get_name() and compute_loss()
    """

    def __init__(self):
        super().__init__()
        self._alpha = 1
        self._loss2 = None

    def compute_loss(self, *args, **kwargs):
        raise NotImplementedError()

    def get_name(self):
        raise NotImplementedError()

    def __mul__(self, alpha):
        assert isinstance(alpha, (int, float))
        res = copy(self)
        res._alpha = alpha
        return res

    __rmul__ = __mul__  # same

    def __add__(self, loss2):
        assert isinstance(loss2, MultiLoss)
        res = cur = copy(self)
        # find the end of the chain
        while cur._loss2 is not None:
            cur = cur._loss2
        cur._loss2 = loss2
        return res

    def __repr__(self):
        name = self.get_name()
        if self._alpha != 1:
            name = f"{self._alpha:g}*{name}"
        if self._loss2:
            name = f"{name} + {self._loss2}"
        return name

    def forward(self, *args, **kwargs):
        loss = self.compute_loss(*args, **kwargs)
        if isinstance(loss, tuple):
            loss, details = loss
        elif loss.ndim == 0:
            details = {self.get_name(): float(loss)}
        else:
            details = {}
        loss = loss * self._alpha

        if self._loss2:
            loss2, details2 = self._loss2(*args, **kwargs)
            loss = loss + loss2
            details |= details2

        return loss, details

class SSIM(nn.Module):
    """Layer to compute the SSIM loss between a pair of images"""

    def __init__(self):
        super(SSIM, self).__init__()
        self.mu_x_pool = nn.AvgPool2d(3, 1)
        self.mu_y_pool = nn.AvgPool2d(3, 1)
        self.sig_x_pool = nn.AvgPool2d(3, 1)
        self.sig_y_pool = nn.AvgPool2d(3, 1)
        self.sig_xy_pool = nn.AvgPool2d(3, 1)

        self.refl = nn.ReflectionPad2d(1)

        self.C1 = 0.01**2
        self.C2 = 0.03**2

    def forward(self, x, y):
        x = self.refl(x)


SSIM

        return name

    def forward(self, *args, **kwargs):
        loss = self.compute_loss(*args, **kwargs)
        if isinstance(loss, tuple):
            loss, details = loss
        elif loss.ndim == 0:
            details = {self.get_name(): float(loss)}
        else:
            details = {}
        loss = loss * self._alpha

        if self._loss2:
            loss2, details2 = self._loss2(*args, **kwargs)
            loss = loss + loss2
            details |= details2

        return loss, details

class SSIM(nn.Module):
    """Layer to compute the SSIM loss between a pair of images"""

    def __init__(self):
        super(SSIM, self).__init__()
        self.mu_x_pool = nn.AvgPool2d(3, 1)
        self.mu_y_pool = nn.AvgPool2d(3, 1)
        self.sig_x_pool = nn.AvgPool2d(3, 1)
        self.sig_y_pool = nn.AvgPool2d(3, 1)
        self.sig_xy_pool = nn.AvgPool2d(3, 1)

        self.refl = nn.ReflectionPad2d(1)

        self.C1 = 0.01**2
        self.C2 = 0.03**2

    def forward(self, x, y):
        x = self.refl(x)
        y = self.refl(y)

        mu_x = self.mu_x_pool(x)
        mu_y = self.mu_y_pool(y)

        sigma_x = self.sig_x_pool(x**2) - mu_x**2
        sigma_y = self.sig_y_pool(y**2) - mu_y**2
        sigma_xy = self.sig_xy_pool(x * y) - mu_x * mu_y

        SSIM_n = (2 * mu_x * mu_y + self.C1) * (2 * sigma_xy + self.C2)
        SSIM_d = (mu_x**2 + mu_y**2 + self.C1) * (sigma_x + sigma_y + self.C2)

        return torch.clamp((1 - SSIM_n / SSIM_d) / 2, 0, 1)

class RGBLoss(Criterion, MultiLoss):
    def __init__(self, criterion):
        super().__init__(criterion)
        self.ssim = SSIM()

    def img_loss(self, a, b):
        return self.criterion(a, b)

    def compute_loss(self, gts, preds, **kw):
        gt_rgbs = [gt["img"].permute(0, 2, 3, 1) for gt in gts]
        pred_rgbs = [pred["rgb"] for pred in preds]
        ls = [
            self.img_loss(pred_rgb, gt_rgb)
            for pred_rgb, gt_rgb in zip(pred_rgbs, gt_rgbs)
        ]
        details = {}
        self_name = type(self).__name__
        for i, l in enumerate(ls):


RGBLoss

        self.C1 = 0.01**2
        self.C2 = 0.03**2

    def forward(self, x, y):
        x = self.refl(x)
        y = self.refl(y)

        mu_x = self.mu_x_pool(x)
        mu_y = self.mu_y_pool(y)

        sigma_x = self.sig_x_pool(x**2) - mu_x**2
        sigma_y = self.sig_y_pool(y**2) - mu_y**2
        sigma_xy = self.sig_xy_pool(x * y) - mu_x * mu_y

        SSIM_n = (2 * mu_x * mu_y + self.C1) * (2 * sigma_xy + self.C2)
        SSIM_d = (mu_x**2 + mu_y**2 + self.C1) * (sigma_x + sigma_y + self.C2)

        return torch.clamp((1 - SSIM_n / SSIM_d) / 2, 0, 1)

class RGBLoss(Criterion, MultiLoss):
    def __init__(self, criterion):
        super().__init__(criterion)
        self.ssim = SSIM()

    def img_loss(self, a, b):
        return self.criterion(a, b)

    def compute_loss(self, gts, preds, **kw):
        gt_rgbs = [gt["img"].permute(0, 2, 3, 1) for gt in gts]
        pred_rgbs = [pred["rgb"] for pred in preds]
        ls = [
            self.img_loss(pred_rgb, gt_rgb)
            for pred_rgb, gt_rgb in zip(pred_rgbs, gt_rgbs)
        ]
        details = {}
        self_name = type(self).__name__
        for i, l in enumerate(ls):
            details[self_name + f"_rgb/{i+1}"] = float(l)
            details[f"pred_rgb_{i+1}"] = pred_rgbs[i]
        rgb_loss = sum(ls) / len(ls)
        return rgb_loss, details

class SMPLLoss(Criterion, MultiLoss):
    def __init__(self, criterion):
        super().__init__(criterion)
        scale = 0.01
        self.alpha_msk = 100.0 * scale
        self.alpha_bce = 10.0 * scale
        self.alpha_rotmat = 100.0 * scale
        self.alpha_shape = 10.0 * scale
        self.alpha_transl = 100.0 * scale
        self.alpha_j3d = 100.0 * scale
        self.alpha_v3d = 100.0 * scale
        self.alpha_j2d = 1.0 * scale
        self.alpha_v2d = 1.0 * scale

        # SMPL layer
        person_center = 'head'
        dict_smpl_layer = {
            'neutral': {


SMPLLoss

    def img_loss(self, a, b):
        return self.criterion(a, b)

    def compute_loss(self, gts, preds, **kw):
        gt_rgbs = [gt["img"].permute(0, 2, 3, 1) for gt in gts]
        pred_rgbs = [pred["rgb"] for pred in preds]
        ls = [
            self.img_loss(pred_rgb, gt_rgb)
            for pred_rgb, gt_rgb in zip(pred_rgbs, gt_rgbs)
        ]
        details = {}
        self_name = type(self).__name__
        for i, l in enumerate(ls):
            details[self_name + f"_rgb/{i+1}"] = float(l)
            details[f"pred_rgb_{i+1}"] = pred_rgbs[i]
        rgb_loss = sum(ls) / len(ls)
        return rgb_loss, details

class SMPLLoss(Criterion, MultiLoss):
    def __init__(self, criterion):
        super().__init__(criterion)
        scale = 0.01
        self.alpha_msk = 100.0 * scale
        self.alpha_bce = 10.0 * scale
        self.alpha_rotmat = 100.0 * scale
        self.alpha_shape = 10.0 * scale
        self.alpha_transl = 100.0 * scale
        self.alpha_j3d = 100.0 * scale
        self.alpha_v3d = 100.0 * scale
        self.alpha_j2d = 1.0 * scale
        self.alpha_v2d = 1.0 * scale

        # SMPL layer
        person_center = 'head'
        dict_smpl_layer = {
            'neutral': {
                10: SMPL_Layer(type='smplx', gender='neutral', num_betas=10, kid=False, person_center=person_center),
                11: SMPL_Layer(type='smplx', gender='neutral', num_betas=11, kid=False, person_center=person_center),
                }
            }
        _moduleDict = []
        for k, _smpl_layer in dict_smpl_layer.items():
            for x, y in _smpl_layer.items():
                _moduleDict.append([f"{k}_{x}", deepcopy(y)])
        self.smpl_layer = nn.ModuleDict(_moduleDict)

    def get_name(self):
        return "SMPLLoss"

    def mask_loss(self, gts, preds, masks, ret_pred=False):
        gt_msks = [gt["msk_mhmr"].unsqueeze(-1) for gt in gts]
        pred_msks = [pred["msk"] for pred in preds]
        ls = [
            F.binary_cross_entropy(p[m], g[m])
            for p, g, m in zip(pred_msks, gt_msks, masks)
        ]
        details = {}
        self_name = self.get_name()
        for i, l in enumerate(ls):
            details[self_name + f"_msk/{i+1}"] = float(l)
            if ret_pred:
                details[f"pred_msk_{i+1}"] = pred_msks[i]
        bce = sum(ls) / len(ls)
        return bce, details

    def bce(self, gts, preds, masks, ret_pred=False):
        gt_scores = [(gt["smpl_scores"] >= 1).to(int).unsqueeze(-1) for gt in gts]
        pred_scores = [pred["smpl_scores"] for pred in preds]
        ls = [
            _neg_loss(p[m], g[m])
            for p, g, m in zip(pred_scores, gt_scores, masks)
        ]
        details = {}
        self_name = self.get_name()
        for i, l in enumerate(ls):
            details[self_name + f"_scores/{i+1}"] = float(l)
            if ret_pred:
                details[f"pred_smpl_scores_{i+1}"] = pred_scores[i]
        bce = sum(ls) / len(ls)
        return bce, details

    def smpl_param_loss(self, gts, preds, masks, k, ret_pred=False):
        if isinstance(gts[0], dict):
            gts = [gt[k] for gt in gts]
        if isinstance(preds[0], dict):
            preds = [pred[k] for pred in preds]

        ls = [
            self.criterion(p[m], g[m])
            for p, g, m in zip(preds, gts, masks)
        ]
        details = {}
        self_name = self.get_name()
        k_name = k.split('smpl_')[-1] if 'smpl_' in k else k
        for i, l in enumerate(ls):
            details[self_name + f"_{k_name}/{i+1}"] = float(l)
            if ret_pred:
                details[f"pred_{k}_{i+1}"] = preds[i]
        loss = sum(ls) / len(ls)

        return loss, details

    def point3d_loss(self, gts, preds, gt_t_p, pr_t_ps, masks, k, ret_pred=False):
        if isinstance(gts[0], dict):
            gts = [gt[k] for gt in gts]
        if isinstance(preds[0], dict):
            preds = [pred[k] for pred in preds]

        if gt_t_p is None or pr_t_ps is None:
            ls = [
                self.criterion(p[m], g[m])
                for p, g, m in zip(preds, gts, masks)
            ]
        else:
            ls = [
                self.criterion(p[m]-pr_p[m], g[m]-gt_p[m])
                for p, pr_p, g, gt_p, m in zip(preds, pr_t_ps, gts, gt_t_p, masks)
            ]

        details = {}
        self_name = self.get_name()
        k_name = k.split('smpl_')[-1] if 'smpl_' in k else k
        if gt_t_p is None or pr_t_ps is None:
            k_name = "c" + k_name
        for i, l in enumerate(ls):
            details[self_name + f"_{k_name}/{i+1}"] = float(l)
            if ret_pred:
                details[f"pred_{k}_{i+1}"] = preds[i]
        loss = sum(ls) / len(ls)

        return loss, details

    def point2d_loss(self, gts, preds, masks, k, shape=None, ret_pred=False):
        if isinstance(gts[0], dict):
            shape = gts[0]['true_shape'][0]
            gts = [gt[k] for gt in gts]
        if isinstance(preds[0], dict):
            preds = [pred[k] for pred in preds]

        valid_mask = [
            ((gt[..., 0] > 0) & (gt[..., 0] < shape[1]) & (gt[..., 1] > 0) & (gt[..., 1] < shape[0]
            )) for gt in gts]

        ls = [
            self.criterion(p[m1.unsqueeze(-1) & m2], g[m1.unsqueeze(-1) & m2])
            for p, g, m1, m2 in zip(preds, gts, masks, valid_mask)
        ]
        details = {}
        self_name = self.get_name()
        k_name = k.split('smpl_')[-1] if 'smpl_' in k else k
        for i, l in enumerate(ls):
            details[self_name + f"_{k_name}/{i+1}"] = float(l)
            if ret_pred:
                details[f"pred_{k}_{i+1}"] = preds[i]
        loss = sum(ls) / len(ls)

        return loss, details

    def compute_loss(self, gts, preds, **kw):
        img_mask_list = [gt["img_mask"] for gt in gts]
        smpl_mask_list = [gt["smpl_mask"] for gt in gts]
        masks_list = [a.unsqueeze(1) & b for a, b in zip(img_mask_list, smpl_mask_list)]

        # Detection loss
        score_loss, score_details = self.bce(gts, preds, img_mask_list, ret_pred=True)

        has_msk = "msk" in preds[0]
        if has_msk:
            msk_loss, msk_details = self.mask_loss(gts, preds, img_mask_list, ret_pred=True)

        K = stack_view(gts, 'camera_intrinsics')
        img_mask = stack_view(img_mask_list,'img_mask').unsqueeze(1)
        smpl_mask = stack_view(smpl_mask_list, 'smpl_mask') * img_mask
        idx_h = torch.where(smpl_mask)
        if int(smpl_mask.sum()) == 0:
            total_loss = self.alpha_bce * score_loss
            details = {
                **score_details,
            }
            if has_msk:
                total_loss += self.alpha_msk * msk_loss
                details.update(msk_details)
            return total_loss, details

        # Prediction
        pred_rotmat = stack_view(preds, 'smpl_rotmat')
        pred_rotvec = roma.rotmat_to_rotvec(pred_rotmat[smpl_mask])
        pred_shape = stack_view(preds, 'smpl_shape')
        pred_transl = [pred.pop("smpl_transl") for pred in preds]
        pred_transl = stack_view(pred_transl, 'smpl_transl')
        pred_expression = stack_view(preds, 'smpl_expression')

        smpl_out = self.smpl_layer[f"neutral_{pred_shape.shape[-1]}"](
            pred_rotvec,
            pred_shape[smpl_mask],
            pred_transl[smpl_mask],
            None, None,
            K=K[idx_h[0]],
            expression=pred_expression[smpl_mask])

        pred_smpl = {}
        batch_size = img_mask_list[0].shape[0]
        num_view = len(gts)
        max_humans = smpl_mask.shape[1]

        for k, v in smpl_out.items():
            full_out = torch.zeros(
                num_view * batch_size, max_humans, *v.shape[1:],
                device=v.device, dtype=v.dtype,
            )
            full_out[smpl_mask] = v
            pred_smpl[k] = full_out.chunk(num_view, dim=0)

        # SMPL-X params
        rotmat_loss, rotmat_details = self.smpl_param_loss(gts, preds, masks_list, "smpl_rotmat")
        transl_loss, transl_details = self.smpl_param_loss(gts, pred_smpl['smpl_transl'], masks_list, "smpl_transl")
        shape_dim = min([gts[0]['smpl_shape'].shape[-1], preds[0]['smpl_shape'].shape[-1]])
        gt_shape = [gt['smpl_shape'][...,:shape_dim] for gt in gts]
        pred_shape = [pred['smpl_shape'][...,:shape_dim] for pred in preds]
        shape_loss, shape_details = self.smpl_param_loss(gt_shape, pred_shape, masks_list, "smpl_shape")

        # 3D points
        gt_transl_pelvis= [gt['smpl_transl_pelvis'][..., None, :] for gt in gts]
        pred_transl_pelvis = pred_smpl['smpl_transl_pelvis']
        j3d_loss, j3d_details = self.point3d_loss(
                                    gts,
                                    pred_smpl['smpl_j3d'],
                                    gt_transl_pelvis,
                                    pred_transl_pelvis,
                                    masks_list, "smpl_j3d")
        v3d_loss, v3d_details = self.point3d_loss(
                                    gts,
                                    pred_smpl['smpl_v3d'],
                                    gt_transl_pelvis,
                                    pred_transl_pelvis,
                                    masks_list, "smpl_v3d",
                                    ret_pred=True)
        cj3d_loss, cj3d_details = self.point3d_loss(
                                    gts,
                                    pred_smpl['smpl_j3d'],
                                    None,
                                    None,
                                    masks_list, "smpl_j3d")
        cv3d_loss, cv3d_details = self.point3d_loss(
                                    gts,
                                    pred_smpl['smpl_v3d'],
                                    None,
                                    None,
                                    masks_list, "smpl_v3d")

        # total loss
        total_loss = self.alpha_bce * score_loss +\
                     self.alpha_rotmat * rotmat_loss +\
                     self.alpha_shape * shape_loss +\
                     self.alpha_transl * transl_loss +\
                     self.alpha_j3d * j3d_loss +\
                     self.alpha_v3d * v3d_loss +\
                     self.alpha_j3d * cj3d_loss +\
                     self.alpha_v3d * cv3d_loss

        details = {
            **score_details,
            **rotmat_details,
            **transl_details,
            **shape_details,
            **j3d_details,
            **v3d_details,
            **cj3d_details,
            **cv3d_details,
        }

        if has_msk:
            total_loss += self.alpha_msk * msk_loss
            details.update(msk_details)

        if cv3d_loss < 1.0:
            # 2D reprojection
            j2d_loss, j2d_details = self.point2d_loss(
                                        gts,
                                        pred_smpl['smpl_j2d'],
                                        masks_list, "smpl_j2d")
            v2d_loss, v2d_details = self.point2d_loss(
                                        gts,
                                        pred_smpl['smpl_v2d'],
                                        masks_list, "smpl_v2d")
            total_loss += self.alpha_j2d * j2d_loss +\
                          self.alpha_v2d * v2d_loss
            details.update({
                **j2d_details,
                **v2d_details
            })

        return total_loss, details

class NaiveSMPLLoss(SMPLLoss):
    def __init__(self, criterion):
        super().__init__(criterion)

    def compute_loss(self, gts, preds, **kw):
        img_mask_list = [gt["img_mask"] for gt in gts]
        smpl_mask_list = [gt["smpl_mask"] for gt in gts]
        masks_list = [a.unsqueeze(1) & b for a, b in zip(img_mask_list, smpl_mask_list)]

        # Detection loss
        score_loss, score_details = self.bce(gts, preds, img_mask_list, ret_pred=True)

        # only for inference SMPL model
        K = stack_view(gts, 'camera_intrinsics')
        img_mask = stack_view(img_mask_list,'img_mask').unsqueeze(1)
        smpl_mask = stack_view(smpl_mask_list, 'smpl_mask') * img_mask
        idx_h = torch.where(smpl_mask)
        if int(smpl_mask.sum()) == 0:
            total_loss = self.alpha_bce * score_loss


NaiveSMPLLoss

        if cv3d_loss < 1.0:
            # 2D reprojection
            j2d_loss, j2d_details = self.point2d_loss(
                                        gts,
                                        pred_smpl['smpl_j2d'],
                                        masks_list, "smpl_j2d")
            v2d_loss, v2d_details = self.point2d_loss(
                                        gts,
                                        pred_smpl['smpl_v2d'],
                                        masks_list, "smpl_v2d")
            total_loss += self.alpha_j2d * j2d_loss +\
                          self.alpha_v2d * v2d_loss
            details.update({
                **j2d_details,
                **v2d_details
            })

        return total_loss, details

class NaiveSMPLLoss(SMPLLoss):
    def __init__(self, criterion):
        super().__init__(criterion)

    def compute_loss(self, gts, preds, **kw):
        img_mask_list = [gt["img_mask"] for gt in gts]
        smpl_mask_list = [gt["smpl_mask"] for gt in gts]
        masks_list = [a.unsqueeze(1) & b for a, b in zip(img_mask_list, smpl_mask_list)]

        # Detection loss
        score_loss, score_details = self.bce(gts, preds, img_mask_list, ret_pred=True)

        # only for inference SMPL model
        K = stack_view(gts, 'camera_intrinsics')
        img_mask = stack_view(img_mask_list,'img_mask').unsqueeze(1)
        smpl_mask = stack_view(smpl_mask_list, 'smpl_mask') * img_mask
        idx_h = torch.where(smpl_mask)
        if int(smpl_mask.sum()) == 0:
            total_loss = self.alpha_bce * score_loss
            details = {
                **score_details,
            }
            return total_loss, details

        # Prediction
        pred_rotmat = stack_view(preds, 'smpl_rotmat')
        pred_rotvec = roma.rotmat_to_rotvec(pred_rotmat[smpl_mask])
        pred_shape = stack_view(preds, 'smpl_shape')
        pred_transl = [pred.pop("smpl_transl") for pred in preds]
        pred_transl = stack_view(pred_transl, 'smpl_transl')
        pred_expression = stack_view(preds, 'smpl_expression')

        # Neutral for MHMR
        K_mhmr = stack_view(gts, 'K_mhmr')
        mhmr_img_res = gts[0]["img_mhmr"].shape[-1]
        # fine head uv
        pred_loc = stack_view(preds, 'smpl_loc')
        # Distance
        dist = pred_transl[smpl_mask][:, 0].unsqueeze(-1)
        dist = to_euclidean_dist(mhmr_img_res, dist, K_mhmr[idx_h[0]])  # use K GT
        smpl_out = self.smpl_layer[f"neutral_{pred_shape.shape[-1]}"](
            pred_rotvec,
            pred_shape[smpl_mask],
            None,
            pred_loc[smpl_mask],
            dist,
            K=K_mhmr[idx_h[0]],
            expression=pred_expression[smpl_mask],
            K_to_proj=K[idx_h[0]], # if use K of CUT3R for projection
            )

        pred_smpl = {}
        batch_size = img_mask_list[0].shape[0]
        num_view = len(gts)
        max_humans = smpl_mask.shape[1]

        for k, v in smpl_out.items():
            full_out = torch.zeros(
                num_view * batch_size, max_humans, *v.shape[1:],
                device=v.device, dtype=v.dtype,
            )
            full_out[smpl_mask] = v
            pred_smpl[k] = full_out.chunk(num_view, dim=0)

        # SMPL-X params
        rotmat_loss, rotmat_details = self.smpl_param_loss(gts, preds, masks_list, "smpl_rotmat")
        transl_loss, transl_details = self.smpl_param_loss(gts, pred_smpl['smpl_transl'], masks_list, "smpl_transl")
        shape_dim = min([gts[0]['smpl_shape'].shape[-1], preds[0]['smpl_shape'].shape[-1]])
        gt_shape = [gt['smpl_shape'][...,:shape_dim] for gt in gts]
        pred_shape = [pred['smpl_shape'][...,:shape_dim] for pred in preds]
        shape_loss, shape_details = self.smpl_param_loss(gt_shape, pred_shape, masks_list, "smpl_shape")

        # 3D points
        gt_transl_pelvis= [gt['smpl_transl_pelvis'][..., None, :] for gt in gts]
        pred_transl_pelvis = pred_smpl['smpl_transl_pelvis']
        j3d_loss, j3d_details = self.point3d_loss(
                                    gts,
                                    pred_smpl['smpl_j3d'],
                                    gt_transl_pelvis,
                                    pred_transl_pelvis,
                                    masks_list, "smpl_j3d")
        v3d_loss, v3d_details = self.point3d_loss(
                                    gts,
                                    pred_smpl['smpl_v3d'],
                                    gt_transl_pelvis,
                                    pred_transl_pelvis,
                                    masks_list, "smpl_v3d",
                                    ret_pred=True)
        cj3d_loss, cj3d_details = self.point3d_loss(
                                    gts,
                                    pred_smpl['smpl_j3d'],
                                    None,
                                    None,
                                    masks_list, "smpl_j3d")
        cv3d_loss, cv3d_details = self.point3d_loss(
                                    gts,
                                    pred_smpl['smpl_v3d'],
                                    None,
                                    None,
                                    masks_list, "smpl_v3d")

        # total loss
        total_loss = self.alpha_bce * score_loss +\
                     self.alpha_rotmat * rotmat_loss +\
                     self.alpha_shape * shape_loss +\
                     self.alpha_transl * transl_loss +\
                     self.alpha_j3d * j3d_loss +\
                     self.alpha_v3d * v3d_loss +\
                     self.alpha_j3d * cj3d_loss +\
                     self.alpha_v3d * cv3d_loss

        details = {
            **score_details,
            **rotmat_details,
            **transl_details,
            **shape_details,
            **j3d_details,
            **v3d_details,
            **cj3d_details,
            **cv3d_details,
        }

        if cv3d_loss < 1.0:
            # 2D reprojection
            j2d_loss, j2d_details = self.point2d_loss(
                                        gts,
                                        pred_smpl['smpl_j2d'],
                                        masks_list, "smpl_j2d")
            v2d_loss, v2d_details = self.point2d_loss(
                                        gts,
                                        pred_smpl['smpl_v2d'],
                                        masks_list, "smpl_v2d")
            total_loss += self.alpha_j2d * j2d_loss +\
                          self.alpha_v2d * v2d_loss
            details.update({
                **j2d_details,
                **v2d_details
            })

        return total_loss, details

class DepthScaleShiftInvLoss(BaseCriterion):
    """scale and shift invariant loss"""

    def __init__(self, reduction="none"):
        super().__init__(reduction)

    def forward(self, pred, gt, mask):
        assert pred.shape == gt.shape and pred.ndim == 3, f"Bad shape = {pred.shape}"
        dist = self.distance(pred, gt, mask)
        # assert dist.ndim == a.ndim - 1  # one dimension less
        if self.reduction == "none":
            return dist
        if self.reduction == "sum":
            return dist.sum()
        if self.reduction == "mean":
            return dist.mean() if dist.numel() > 0 else dist.new_zeros(())
        raise ValueError(f"bad {self.reduction=} mode")

    def normalize(self, x, mask):


DepthScaleShiftInvLoss

        if cv3d_loss < 1.0:
            # 2D reprojection
            j2d_loss, j2d_details = self.point2d_loss(
                                        gts,
                                        pred_smpl['smpl_j2d'],
                                        masks_list, "smpl_j2d")
            v2d_loss, v2d_details = self.point2d_loss(
                                        gts,
                                        pred_smpl['smpl_v2d'],
                                        masks_list, "smpl_v2d")
            total_loss += self.alpha_j2d * j2d_loss +\
                          self.alpha_v2d * v2d_loss
            details.update({
                **j2d_details,
                **v2d_details
            })

        return total_loss, details

class DepthScaleShiftInvLoss(BaseCriterion):
    """scale and shift invariant loss"""

    def __init__(self, reduction="none"):
        super().__init__(reduction)

    def forward(self, pred, gt, mask):
        assert pred.shape == gt.shape and pred.ndim == 3, f"Bad shape = {pred.shape}"
        dist = self.distance(pred, gt, mask)
        # assert dist.ndim == a.ndim - 1  # one dimension less
        if self.reduction == "none":
            return dist
        if self.reduction == "sum":
            return dist.sum()
        if self.reduction == "mean":
            return dist.mean() if dist.numel() > 0 else dist.new_zeros(())
        raise ValueError(f"bad {self.reduction=} mode")

    def normalize(self, x, mask):
        x_valid = x[mask]
        splits = mask.sum(dim=(1, 2)).tolist()
        x_valid_list = torch.split(x_valid, splits)
        shift = [x.mean() for x in x_valid_list]
        x_valid_centered = [x - m for x, m in zip(x_valid_list, shift)]
        scale = [x.abs().mean() for x in x_valid_centered]
        scale = torch.stack(scale)
        shift = torch.stack(shift)
        x = (x - shift.view(-1, 1, 1)) / scale.view(-1, 1, 1).clamp(min=1e-6)
        return x

    def distance(self, pred, gt, mask):
        pred = self.normalize(pred, mask)
        gt = self.normalize(gt, mask)
        return torch.abs((pred - gt)[mask])

class ScaleInvLoss(BaseCriterion):
    """scale invariant loss"""

    def __init__(self, reduction="none"):
        super().__init__(reduction)

    def forward(self, pred, gt, mask):
        assert pred.shape == gt.shape and pred.ndim == 4, f"Bad shape = {pred.shape}"
        dist = self.distance(pred, gt, mask)
        # assert dist.ndim == a.ndim - 1  # one dimension less
        if self.reduction == "none":
            return dist
        if self.reduction == "sum":
            return dist.sum()
        if self.reduction == "mean":
            return dist.mean() if dist.numel() > 0 else dist.new_zeros(())
        raise ValueError(f"bad {self.reduction=} mode")


ScaleInvLoss

        raise ValueError(f"bad {self.reduction=} mode")

    def normalize(self, x, mask):
        x_valid = x[mask]
        splits = mask.sum(dim=(1, 2)).tolist()
        x_valid_list = torch.split(x_valid, splits)
        shift = [x.mean() for x in x_valid_list]
        x_valid_centered = [x - m for x, m in zip(x_valid_list, shift)]
        scale = [x.abs().mean() for x in x_valid_centered]
        scale = torch.stack(scale)
        shift = torch.stack(shift)
        x = (x - shift.view(-1, 1, 1)) / scale.view(-1, 1, 1).clamp(min=1e-6)
        return x

    def distance(self, pred, gt, mask):
        pred = self.normalize(pred, mask)
        gt = self.normalize(gt, mask)
        return torch.abs((pred - gt)[mask])

class ScaleInvLoss(BaseCriterion):
    """scale invariant loss"""

    def __init__(self, reduction="none"):
        super().__init__(reduction)

    def forward(self, pred, gt, mask):
        assert pred.shape == gt.shape and pred.ndim == 4, f"Bad shape = {pred.shape}"
        dist = self.distance(pred, gt, mask)
        # assert dist.ndim == a.ndim - 1  # one dimension less
        if self.reduction == "none":
            return dist
        if self.reduction == "sum":
            return dist.sum()
        if self.reduction == "mean":
            return dist.mean() if dist.numel() > 0 else dist.new_zeros(())
        raise ValueError(f"bad {self.reduction=} mode")

    def distance(self, pred, gt, mask):
        pred_norm_factor = (torch.norm(pred, dim=-1) * mask).sum(dim=(1, 2)) / mask.sum(
            dim=(1, 2)
        ).clamp(min=1e-6)
        gt_norm_factor = (torch.norm(gt, dim=-1) * mask).sum(dim=(1, 2)) / mask.sum(
            dim=(1, 2)
        ).clamp(min=1e-6)
        pred = pred / pred_norm_factor.view(-1, 1, 1, 1).clamp(min=1e-6)
        gt = gt / gt_norm_factor.view(-1, 1, 1, 1).clamp(min=1e-6)
        return torch.norm(pred - gt, dim=-1)[mask]

class Regr3DPose(Criterion, MultiLoss):
    """Ensure that all 3D points are correct.
    Asymmetric loss: view1 is supposed to be the anchor.

    P1 = RT1 @ D1
    P2 = RT2 @ D2
    loss1 = (I @ pred_D1) - (RT1^-1 @ RT1 @ D1)
    loss2 = (RT21 @ pred_D2) - (RT1^-1 @ P2)
          = (RT21 @ pred_D2) - (RT1^-1 @ RT2 @ D2)
    """

    def __init__(
        self,
        criterion,
        norm_mode="?avg_dis",
        gt_scale=False,
        sky_loss_value=2,
        max_metric_scale=False,