import os
import fire
import gradio as gr
from PIL import Image
from functools import partial
import argparse
import sys
import torch

if os.getenv('SYSTEM') == 'spaces':
    os.system('pip install --global-option="--no-networks" git+https://github.com/NVlabs/tiny-cuda-nn/#subdirectory=bindings/torch')
    # os.system('pip install fvcore iopath')
    # os.system('pip install "git+https://github.com/facebookresearch/pytorch3d.git"')
    
import cv2
import time
import numpy as np
import trimesh
from segment_anything import build_sam, SamPredictor

import random
from pytorch3d import transforms
import torchvision
import torch.distributed as dist
import nvdiffrast.torch as dr
from video3d.model_ddp import Unsup3DDDP, forward_to_matrix
from video3d.trainer_few_shot import Fewshot_Trainer
from video3d.trainer_ddp import TrainerDDP
from video3d import setup_runtime
from video3d.render.mesh import make_mesh
from video3d.utils.skinning_v4 import estimate_bones, skinning, euler_angles_to_matrix
from video3d.utils.misc import save_obj
from video3d.render import util
import matplotlib.pyplot as plt
from pytorch3d import utils, renderer, transforms, structures, io
from video3d.render.render import render_mesh
from video3d.render.material import texture as material_texture


_TITLE = '''Learning the 3D Fauna of the Web'''
_DESCRIPTION = '''
<div>
Reconstruct any quadruped animal from one image.
</div>
<div>
The demo only contains the 3D reconstruction part.
</div>
'''
_GPU_ID = 0

if not hasattr(Image, 'Resampling'):
    Image.Resampling = Image


def sam_init():
    sam_checkpoint = os.path.join(os.path.dirname(__file__), "sam_pt", "sam_vit_h_4b8939.pth")
    model_type = "vit_h"
    # sam = sam_model_registry[model_type](checkpoint=sam_checkpoint).to(device=f"cuda:{_GPU_ID}")
    # predictor = SamPredictor(sam)
    predictor = SamPredictor(build_sam(checkpoint=sam_checkpoint).to("cuda"))
    return predictor


def sam_segment(predictor, input_image, *bbox_coords):
    bbox = np.array(bbox_coords)
    image = np.asarray(input_image)

    start_time = time.time()
    predictor.set_image(image)

    masks_bbox, scores_bbox, logits_bbox = predictor.predict(
        box=bbox,
        multimask_output=True
    )

    print(f"SAM Time: {time.time() - start_time:.3f}s")
    out_image = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
    out_image[:, :, :3] = image
    out_image_bbox = out_image.copy()
    out_image_bbox[:, :, 3] = masks_bbox[-1].astype(np.uint8) * 255
    torch.cuda.empty_cache()
    # return Image.fromarray(out_image_bbox, mode='RGB')

    x_nonzero = np.nonzero(masks_bbox[-1].astype(np.uint8).sum(axis=0))
    y_nonzero = np.nonzero(masks_bbox[-1].astype(np.uint8).sum(axis=1))
    x_min = int(x_nonzero[0].min())
    y_min = int(y_nonzero[0].min())
    x_max = int(x_nonzero[0].max())
    y_max = int(y_nonzero[0].max())

    out_image_bbox = out_image_bbox[y_min:y_max, x_min:x_max]

    out_image_valid = np.concatenate([
        out_image_bbox[:, :, :3],
        np.ones_like(out_image_bbox[:, :, 3:]) * 255
    ], axis=-1)

    return Image.fromarray(out_image_valid, mode='RGBA')


def expand2square(pil_img, background_color):
    width, height = pil_img.size
    if width == height:
        return pil_img
    elif width > height:
        result = Image.new(pil_img.mode, (width, width), background_color)
        result.paste(pil_img, (0, (width - height) // 2))
        return result
    else:
        result = Image.new(pil_img.mode, (height, height), background_color)
        result.paste(pil_img, ((height - width) // 2, 0))
        return result


def preprocess(predictor, input_image, chk_group=None, segment=False):
    RES = 1024
    input_image.thumbnail([RES, RES], Image.Resampling.LANCZOS)
    if chk_group is not None:
        segment = "Use SAM to center animal" in chk_group
    if segment:
        image_rem = input_image.convert('RGB')
        arr = np.asarray(image_rem)[:,:,-1]
        x_nonzero = np.nonzero(arr.sum(axis=0))
        y_nonzero = np.nonzero(arr.sum(axis=1))
        x_min = int(x_nonzero[0].min())
        y_min = int(y_nonzero[0].min())
        x_max = int(x_nonzero[0].max())
        y_max = int(y_nonzero[0].max())
        input_image = sam_segment(predictor, input_image.convert('RGB'), x_min, y_min, x_max, y_max)
    
    input_image = expand2square(input_image, (0, 0, 0, 255))
    return input_image, input_image.resize((256, 256), Image.Resampling.LANCZOS)


def save_images(images, mask_pred, mode="transparent"):
    img = images[0]
    mask = mask_pred[0]
    img = img.clamp(0, 1)
    if mask is not None:
        mask = mask.clamp(0, 1)
        if mode == "white":
            img = img * mask + 1 * (1 - mask)
        elif mode == "black":
            img = img * mask + 0 * (1 - mask)
        else:
            img = torch.cat([img, mask[0:1]], 0)
    
    img = img.permute(1, 2, 0).cpu().numpy()
    img = Image.fromarray(np.uint8(img * 255))
    return img


def get_bank_embedding(rgb, memory_bank_keys, memory_bank, model, memory_bank_topk=10, memory_bank_dim=128):
    images = rgb
    batch_size, num_frames, _, h0, w0 = images.shape
    images = images.reshape(batch_size*num_frames, *images.shape[2:])  # 0~1
    images_in = images * 2 - 1  # rescale to (-1, 1) for DINO
    
    x = images_in
    with torch.no_grad():
        b, c, h, w = x.shape
        model.netInstance.netEncoder._feats = []
        model.netInstance.netEncoder._register_hooks([11], 'key')
        #self._register_hooks([11], 'token')
        x = model.netInstance.netEncoder.ViT.prepare_tokens(x)
        #x = self.ViT.prepare_tokens_with_masks(x)
        
        for blk in model.netInstance.netEncoder.ViT.blocks:
            x = blk(x)
        out = model.netInstance.netEncoder.ViT.norm(x)
        model.netInstance.netEncoder._unregister_hooks()

        ph, pw = h // model.netInstance.netEncoder.patch_size, w // model.netInstance.netEncoder.patch_size
        patch_out = out[:, 1:]  # first is class token
        patch_out = patch_out.reshape(b, ph, pw, model.netInstance.netEncoder.vit_feat_dim).permute(0, 3, 1, 2)

        patch_key = model.netInstance.netEncoder._feats[0][:,:,1:]  # B, num_heads, num_patches, dim
        patch_key = patch_key.permute(0, 1, 3, 2).reshape(b, model.netInstance.netEncoder.vit_feat_dim, ph, pw)

        global_feat = out[:, 0]
    
    batch_features = global_feat

    batch_size = batch_features.shape[0]
        
    query = torch.nn.functional.normalize(batch_features.unsqueeze(1), dim=-1)      # [B, 1, d_k]
    key = torch.nn.functional.normalize(memory_bank_keys, dim=-1)              # [size, d_k]
    key = key.transpose(1, 0).unsqueeze(0).repeat(batch_size, 1, 1).to(query.device)             # [B, d_k, size]

    cos_dist = torch.bmm(query, key).squeeze(1)         # [B, size], larger the more similar
    rank_idx = torch.sort(cos_dist, dim=-1, descending=True)[1][:, :memory_bank_topk] # [B, k]
    value = memory_bank.unsqueeze(0).repeat(batch_size, 1, 1).to(query.device)                         # [B, size, d_v]

    out = torch.gather(value, dim=1, index=rank_idx[..., None].repeat(1, 1, memory_bank_dim))  # [B, k, d_v]

    weights = torch.gather(cos_dist, dim=-1, index=rank_idx)    # [B, k]
    weights = torch.nn.functional.normalize(weights, p=1.0, dim=-1).unsqueeze(-1).repeat(1, 1, memory_bank_dim)    # [B, k, d_v] weights have been normalized

    out = weights * out
    out = torch.sum(out, dim=1)
    
    batch_mean_out = torch.mean(out, dim=0)

    weight_aux = {
        'weights': weights[:, :, 0], # [B, k], weights from large to small
        'pick_idx': rank_idx, # [B, k]
    }

    batch_embedding = batch_mean_out 
    embeddings = out
    weights = weight_aux

    bank_embedding_model_input = [batch_embedding, embeddings, weights]

    return bank_embedding_model_input


class FixedDirectionLight(torch.nn.Module):
    def __init__(self, direction, amb, diff):
        super(FixedDirectionLight, self).__init__()
        self.light_dir = direction
        self.amb = amb
        self.diff = diff
        self.is_hacking = not (isinstance(self.amb, float)
                               or isinstance(self.amb, int))

    def forward(self, feat):
        batch_size = feat.shape[0]
        if self.is_hacking:
            return torch.concat([self.light_dir, self.amb, self.diff], -1)
        else:
            return torch.concat([self.light_dir, torch.FloatTensor([self.amb, self.diff]).to(self.light_dir.device)], -1).expand(batch_size, -1)

    def shade(self, feat, kd, normal):
        light_params = self.forward(feat)
        light_dir = light_params[..., :3][:, None, None, :]
        int_amb = light_params[..., 3:4][:, None, None, :]
        int_diff = light_params[..., 4:5][:, None, None, :]
        shading = (int_amb + int_diff *
                   torch.clamp(util.dot(light_dir, normal), min=0.0))
        shaded = shading * kd
        return shaded, shading


def render_bones(mvp, bones_pred, size=(256, 256)):
    bone_world4 = torch.concat([bones_pred, torch.ones_like(bones_pred[..., :1]).to(bones_pred.device)], dim=-1)
    b, f, num_bones = bone_world4.shape[:3]
    bones_clip4 = (bone_world4.view(b, f, num_bones*2, 1, 4) @ mvp.transpose(-1, -2).reshape(b, f, 1, 4, 4)).view(b, f, num_bones, 2, 4)
    bones_uv = bones_clip4[..., :2] / bones_clip4[..., 3:4]  # b, f, num_bones, 2, 2
    dpi = 32
    fx, fy = size[1] // dpi, size[0] // dpi

    rendered = []
    for b_idx in range(b):
        for f_idx in range(f):
            frame_bones_uv = bones_uv[b_idx, f_idx].cpu().numpy()
            fig = plt.figure(figsize=(fx, fy), dpi=dpi, frameon=False)
            ax = plt.Axes(fig, [0., 0., 1., 1.])
            ax.set_axis_off()
            for bone in frame_bones_uv:
                ax.plot(bone[:, 0], bone[:, 1], marker='o', linewidth=8, markersize=20)
            ax.set_xlim(-1, 1)
            ax.set_ylim(-1, 1)
            ax.invert_yaxis()
            # Convert to image
            fig.add_axes(ax)
            fig.canvas.draw_idle()
            image = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
            w, h = fig.canvas.get_width_height()
            image.resize(h, w, 3)
            rendered += [image / 255.]
    return torch.from_numpy(np.stack(rendered, 0).transpose(0, 3, 1, 2)).to(bones_pred.device)

def add_mesh_color(mesh, color):
    verts = mesh.verts_padded()
    color = torch.FloatTensor(color).to(verts.device).view(1,1,3) / 255
    mesh.textures = renderer.TexturesVertex(verts_features=verts*0+color)
    return mesh

def create_sphere(position, scale, device, color=[139, 149, 173]):
    mesh = utils.ico_sphere(2).to(device)
    mesh = mesh.extend(position.shape[0])

    # scale and offset
    mesh = mesh.update_padded(mesh.verts_padded() * scale + position[:, None])

    mesh = add_mesh_color(mesh, color)

    return mesh

def estimate_bone_rotation(b):
    """
    (0, 0, 1) = matmul(R^(-1), b)

    assumes x, y is a symmetry plane

    returns R
    """
    b = b / torch.norm(b, dim=-1, keepdim=True)

    n = torch.FloatTensor([[1, 0, 0]]).to(b.device)
    n = n.expand_as(b)
    v = torch.cross(b, n, dim=-1)

    R = torch.stack([n, v, b], dim=-1).transpose(-2, -1)

    return R

def estimate_vector_rotation(vector_a, vector_b):
    """
    vector_a = matmul(R, vector_b)

    returns R

    https://math.stackexchange.com/questions/180418/calculate-rotation-matrix-to-align-vector-a-to-vector-b-in-3d
    """
    vector_a = vector_a / torch.norm(vector_a, dim=-1, keepdim=True)
    vector_b = vector_b / torch.norm(vector_b, dim=-1, keepdim=True)

    v = torch.cross(vector_a, vector_b, dim=-1)
    c = torch.sum(vector_a * vector_b, dim=-1)

    skew = torch.stack([
        torch.stack([torch.zeros_like(v[..., 0]), -v[..., 2], v[..., 1]], dim=-1),
        torch.stack([v[..., 2], torch.zeros_like(v[..., 0]), -v[..., 0]], dim=-1),
        torch.stack([-v[..., 1], v[..., 0], torch.zeros_like(v[..., 0])], dim=-1)],
        dim=-1)

    R = torch.eye(3, device=vector_a.device)[None] + skew + torch.matmul(skew, skew) / (1  + c[..., None, None])

    return R

def create_elipsoid(bone, scale=0.05, color=[139, 149, 173], generic_rotation_estim=True):
    length = torch.norm(bone[:, 0] - bone[:, 1], dim=-1)

    mesh = utils.ico_sphere(2).to(bone.device)
    mesh = mesh.extend(bone.shape[0])
    # scale x, y
    verts = mesh.verts_padded() * torch.FloatTensor([scale, scale, 1]).to(bone.device)
    # stretch along z axis, set the start to origin
    verts[:, :, 2] = verts[:, :, 2] * length[:, None] * 0.5 + length[:, None] * 0.5

    bone_vector = bone[:, 1] - bone[:, 0]
    z_vector = torch.FloatTensor([[0, 0, 1]]).to(bone.device)
    z_vector = z_vector.expand_as(bone_vector)
    if generic_rotation_estim:
        rot = estimate_vector_rotation(z_vector, bone_vector)
    else:
        rot = estimate_bone_rotation(bone_vector)
    tsf = transforms.Rotate(rot, device=bone.device)
    tsf = tsf.compose(transforms.Translate(bone[:, 0], device=bone.device))
    verts = tsf.transform_points(verts)

    mesh = mesh.update_padded(verts)

    mesh = add_mesh_color(mesh, color)

    return mesh

def convert_textures_vertex_to_textures_uv(meshes: structures.Meshes, color1, color2) -> renderer.TexturesUV:
    """
    Convert a TexturesVertex object to a TexturesUV object.
    """
    color1 = torch.Tensor(color1).to(meshes.device).view(1, 1, 3) / 255
    color2 = torch.Tensor(color2).to(meshes.device).view(1, 1, 3) / 255
    textures_vertex = meshes.textures
    assert isinstance(textures_vertex, renderer.TexturesVertex), "Input meshes must have TexturesVertex"
    verts_rgb = textures_vertex.verts_features_padded()
    faces_uvs = meshes.faces_padded()
    batch_size = verts_rgb.shape[0]
    maps = torch.zeros(batch_size, 128, 128, 3, device=verts_rgb.device)
    maps[:, :, :64, :] = color1
    maps[:, :, 64:, :] = color2
    is_first = (verts_rgb == color1)[..., 0]
    verts_uvs = torch.zeros(batch_size, verts_rgb.shape[1], 2, device=verts_rgb.device)
    verts_uvs[is_first] = torch.FloatTensor([0.25, 0.5]).to(verts_rgb.device)
    verts_uvs[~is_first] = torch.FloatTensor([0.75, 0.5]).to(verts_rgb.device)
    textures_uv = renderer.TexturesUV(maps=maps, faces_uvs=faces_uvs, verts_uvs=verts_uvs)
    meshes.textures = textures_uv
    return meshes
    
def create_bones_scene(bones, joint_color=[66, 91, 140], bone_color=[119, 144, 189], show_end_point=False):
    meshes = []
    for bone_i in range(bones.shape[1]):
        # points
        meshes += [create_sphere(bones[:, bone_i, 0], 0.1, bones.device, color=joint_color)]
        if show_end_point:
            meshes += [create_sphere(bones[:, bone_i, 1], 0.1, bones.device, color=joint_color)]

        # connecting ellipsoid
        meshes += [create_elipsoid(bones[:, bone_i], color=bone_color)]

    current_batch_size = bones.shape[0]
    meshes = [structures.join_meshes_as_scene([m[i] for m in meshes]) for i in range(current_batch_size)]
    mesh = structures.join_meshes_as_batch(meshes)

    return mesh


def save_mesh(mesh, file_path):
    obj_file = file_path
    idx = 0
    print("Writing mesh: ", obj_file)
    with open(obj_file, "w") as f:
        # f.write(f"mtllib {fname}.mtl\n")
        f.write("g default\n")

        v_pos = mesh.v_pos[idx].detach().cpu().numpy() if mesh.v_pos is not None else None
        v_nrm = mesh.v_nrm[idx].detach().cpu().numpy() if mesh.v_nrm is not None else None
        v_tex = mesh.v_tex[idx].detach().cpu().numpy() if mesh.v_tex is not None else None

        t_pos_idx = mesh.t_pos_idx[0].detach().cpu().numpy() if mesh.t_pos_idx is not None else None
        t_nrm_idx = mesh.t_nrm_idx[0].detach().cpu().numpy() if mesh.t_nrm_idx is not None else None
        t_tex_idx = mesh.t_tex_idx[0].detach().cpu().numpy() if mesh.t_tex_idx is not None else None

        print("    writing %d vertices" % len(v_pos))
        for v in v_pos:
            f.write('v {} {} {} \n'.format(v[0], v[1], v[2]))

        if v_nrm is not None:
            print("    writing %d normals" % len(v_nrm))
            assert(len(t_pos_idx) == len(t_nrm_idx))
            for v in v_nrm:
                f.write('vn {} {} {}\n'.format(v[2], v[1], v[0]))

        # faces
        f.write("s 1 \n")
        f.write("g pMesh1\n")
        f.write("usemtl defaultMat\n")

        # Write faces
        print("    writing %d faces" % len(t_pos_idx))
        for i in range(len(t_pos_idx)):
            f.write("f ")
            for j in range(3):
                f.write(' %s/%s/%s' % (str(t_pos_idx[i][j]+1), '' if v_tex is None else str(t_tex_idx[i][j]+1), '' if v_nrm is None else str(t_nrm_idx[i][j]+1)))
            f.write("\n")


def process_mesh(shape, name):
    mesh = shape.clone()
    output_glb = f'./{name}.glb'
    output_obj = f'./{name}.obj'

    # save the obj file for download
    save_mesh(mesh, output_obj)

    # save the glb for visualize
    mesh_tri = trimesh.Trimesh(
        vertices=mesh.v_pos[0].detach().cpu().numpy(), 
        faces=mesh.t_pos_idx[0][..., [2,1,0]].detach().cpu().numpy(), 
        process=False,  
        maintain_order=True
    )

    norm_colors = (mesh.v_nrm[0][..., [2,1,0]].detach().cpu().numpy() + 1.0) * 0.5 * 1.8 * 255.0
    norm_colors = np.clip(norm_colors, 0, 255)

    mesh_tri.visual.vertex_colors = norm_colors
    mesh_tri.export(file_obj=output_glb)

    return output_glb, output_obj


def run_pipeline(model_items, cfgs, input_img):
    epoch = 999
    total_iter = 999999
    model = model_items[0]
    memory_bank = model_items[1]
    memory_bank_keys = model_items[2]

    device = f'cuda:{_GPU_ID}'

    input_image = torch.stack([torchvision.transforms.ToTensor()(input_img)], dim=0).to(device)

    with torch.no_grad():
        model.netPrior.eval()
        model.netInstance.eval()
        input_image = torch.nn.functional.interpolate(input_image, size=(256, 256), mode='bilinear', align_corners=False)
        input_image = input_image[:, None, :, :]  # [B=1, F=1, 3, 256, 256]

        bank_embedding = get_bank_embedding(
            input_image, 
            memory_bank_keys, 
            memory_bank, 
            model, 
            memory_bank_topk=cfgs.get("memory_bank_topk", 10),
            memory_bank_dim=128
        )

        prior_shape, dino_pred, classes_vectors = model.netPrior(
            category_name='tmp',
            perturb_sdf=False, 
            total_iter=total_iter,
            is_training=False, 
            class_embedding=bank_embedding
        )
        Instance_out = model.netInstance(
            'tmp', 
            input_image, 
            prior_shape, 
            epoch, 
            dino_features=None, 
            dino_clusters=None, 
            total_iter=total_iter, 
            is_training=False
        )  # frame dim collapsed N=(B*F)
        if len(Instance_out) == 13:
            shape, pose_raw, pose, mvp, w2c, campos, texture_pred, im_features, dino_feat_im_calc, deform, all_arti_params, light, forward_aux = Instance_out
            im_features_map = None
        else:
            shape, pose_raw, pose, mvp, w2c, campos, texture_pred, im_features, dino_feat_im_calc, deform, all_arti_params, light, forward_aux, im_features_map = Instance_out

        class_vector = classes_vectors  # the bank embeddings

        gray_light = FixedDirectionLight(direction=torch.FloatTensor([0, 0, 1]).to(device), amb=0.2, diff=0.7)

        image_pred, mask_pred, _, _, _, shading = model.render(
            shape, texture_pred, mvp, w2c, campos, (256, 256), background=model.background_mode, 
            im_features=im_features, light=gray_light, prior_shape=prior_shape, render_mode='diffuse', 
            render_flow=False, dino_pred=None, im_features_map=im_features_map
        )
        mask_pred = mask_pred.expand_as(image_pred)
        shading = shading.expand_as(image_pred)
        # render bones in pytorch3D style
        posed_bones = forward_aux["posed_bones"].squeeze(1)
        jc, bc = [66, 91, 140], [119, 144, 189]
        bones_meshes = create_bones_scene(posed_bones, joint_color=jc, bone_color=bc, show_end_point=True)
        bones_meshes = convert_textures_vertex_to_textures_uv(bones_meshes, color1=jc, color2=bc)
        nv_meshes = make_mesh(verts=bones_meshes.verts_padded(), faces=bones_meshes.faces_padded()[0:1],
                                uvs=bones_meshes.textures.verts_uvs_padded(), uv_idx=bones_meshes.textures.faces_uvs_padded()[0:1],
                                material=material_texture.Texture2D(bones_meshes.textures.maps_padded()))
        # buffers = render_mesh(dr.RasterizeGLContext(), nv_meshes, mvp, w2c, campos, nv_meshes.material, lgt=gray_light, feat=im_features, dino_pred=None, resolution=(256,256), bsdf="diffuse")
        buffers = render_mesh(dr.RasterizeCudaContext(), nv_meshes, mvp, w2c, campos, nv_meshes.material, lgt=gray_light, feat=im_features, dino_pred=None, resolution=(256,256), bsdf="diffuse")
        
        shaded = buffers["shaded"].permute(0, 3, 1, 2)
        bone_image = shaded[:, :3, :, :]
        bone_mask = shaded[:, 3:, :, :]
        mask_final = mask_pred.logical_or(bone_mask)
        mask_final = mask_final.int()
        image_with_bones = bone_image * bone_mask * 0.5 + (shading * (1 - bone_mask * 0.5) + 0.5 * (mask_final.float() - mask_pred.float()))

        mesh_image = save_images(shading, mask_pred)
        mesh_bones_image = save_images(image_with_bones, mask_final)

        shape_glb, shape_obj = process_mesh(shape, 'reconstructed_shape')
        base_shape_glb, base_shape_obj = process_mesh(prior_shape, 'reconstructed_base_shape')

        return mesh_image, mesh_bones_image, shape_glb, shape_obj, base_shape_glb, base_shape_obj


def run_demo():
    parser = argparse.ArgumentParser()
    parser.add_argument('--gpu', default='0', type=str,
                        help='Specify a GPU device')
    parser.add_argument('--num_workers', default=4, type=int,
                        help='Specify the number of worker threads for data loaders')
    parser.add_argument('--seed', default=0, type=int,
                        help='Specify a random seed')
    parser.add_argument('--config', default='./ckpts/configs.yml',
                        type=str)  # Model config path
    parser.add_argument('--checkpoint_path', default='./ckpts/iter0800000.pth', type=str)

    args = parser.parse_args()

    torch.manual_seed(args.seed)
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '8088'
    dist.init_process_group("gloo", rank=_GPU_ID, world_size=1)
    torch.cuda.set_device(_GPU_ID)
    args.rank = _GPU_ID
    args.world_size = 1
    args.gpu = f'{_GPU_ID}'
    device = f'cuda:{_GPU_ID}'

    resolution = (256, 256)
    batch_size = 1
    model_cfgs = setup_runtime(args)
    bone_y_thresh = 0.4
    body_bone_idx_preset = [3, 6, 6, 3]
    model_cfgs['body_bone_idx_preset'] = body_bone_idx_preset

    model = Unsup3DDDP(model_cfgs)
    # a hack attempt
    model.netPrior.classes_vectors = torch.nn.Parameter(torch.nn.init.uniform_(torch.empty(123, 128), a=-0.05, b=0.05))
    cp = torch.load(args.checkpoint_path, map_location=device)
    model.load_model_state(cp)
    memory_bank_keys = cp['memory_bank_keys']
    memory_bank = cp['memory_bank']

    model.to(device)
    memory_bank.to(device)
    memory_bank_keys.to(device)
    model_items = [
        model,
        memory_bank,
        memory_bank_keys
    ]

    predictor = sam_init()

    custom_theme = gr.themes.Soft(primary_hue="blue").set(
                    button_secondary_background_fill="*neutral_100",
                    button_secondary_background_fill_hover="*neutral_200")
    custom_css = '''#disp_image {
        text-align: center; /* Horizontally center the content */
    }'''

    with gr.Blocks(title=_TITLE, theme=custom_theme, css=custom_css) as demo:
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown('# ' + _TITLE)
        gr.Markdown(_DESCRIPTION)
        with gr.Row(variant='panel'):
            with gr.Column(scale=1):
                input_image = gr.Image(type='pil', image_mode='RGBA', height=256, label='Input image', tool=None)

                example_folder = os.path.join(os.path.dirname(__file__), "./example_images")
                example_fns = [os.path.join(example_folder, example) for example in os.listdir(example_folder)]
                gr.Examples(
                    examples=example_fns,
                    inputs=[input_image],
                    # outputs=[input_image],
                    cache_examples=False,
                    label='Examples (click one of the images below to start)',
                    examples_per_page=30
                )
            with gr.Column(scale=1):
                processed_image = gr.Image(type='pil', label="Processed Image", interactive=False, height=256, tool=None, image_mode='RGB', elem_id="disp_image")
                processed_image_highres = gr.Image(type='pil', image_mode='RGB', visible=False, tool=None)

                with gr.Accordion('Advanced options', open=True):
                    with gr.Row():
                        with gr.Column():
                            input_processing = gr.CheckboxGroup(['Use SAM to center animal'], 
                                                                label='Input Image Preprocessing',
                                                                 info='untick this, if animal is already centered, e.g. in example images')
                        # with gr.Column():
                        #     output_processing = gr.CheckboxGroup(['Background Removal'], label='Output Image Postprocessing', value=[]) 
                    # with gr.Row():
                    #     with gr.Column():
                    #         scale_slider = gr.Slider(1, 5, value=3, step=1,
                    #                                     label='Classifier Free Guidance Scale')
                    #     with gr.Column():
                    #         steps_slider = gr.Slider(15, 100, value=50, step=1,
                    #                                     label='Number of Diffusion Inference Steps')
                    # with gr.Row():
                    #     with gr.Column():
                    #         seed = gr.Number(42, label='Seed')
                    #     with gr.Column():
                    #         crop_size = gr.Number(192, label='Crop size')
                    # crop_size = 192
                run_btn = gr.Button('Reconstruct', variant='primary', interactive=True)
        with gr.Row():
            view_1 = gr.Image(label="Input View Reconstruction", interactive=False, height=256, show_label=True)
            view_2 = gr.Image(label="Input View Reconstruction with Skeleton", interactive=False, height=256, show_label=True)
        with gr.Row():
            shape_1 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  height=512, label="Reconstructed Shape")
            shape_2 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  height=512, label="Reconstructed Base Shape")
            #shape_1_download = gr.File(label="Download Full Reconstructed Model")
        with gr.Row():
            # shape_2 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0],  height=512, label="Bank Base Shape Model")
            shape_1_download = gr.File(label="Download Reconstructed Shape")
            shape_2_download = gr.File(label="Download Reconstructed Base Shape")

        run_btn.click(fn=partial(preprocess, predictor), 
                        inputs=[input_image, input_processing], 
                        outputs=[processed_image_highres, processed_image], queue=True
            ).success(fn=partial(run_pipeline, model_items, model_cfgs), 
                        inputs=[processed_image],
                        outputs=[view_1, view_2, shape_1, shape_1_download, shape_2, shape_2_download]
                        )
        demo.queue().launch(share=True, max_threads=80)
        # _, local_url, share_url = demo.queue().launch(share=True, server_name="0.0.0.0", server_port=23425)
        # print('local_url: ', local_url)


if __name__ == '__main__':
    fire.Fire(run_demo)