Custom mini dataset for TripoSR model traing?

[developer-kliky-app]

Model Configuration look like this @bennyguo @pookiefoof @kostarion @varunjampani

cond_image_size: 512

image_tokenizer_cls: tsr.models.tokenizers.image.DINOSingleImageTokenizer
image_tokenizer:
  pretrained_model_name_or_path: "facebook/dino-vitb16"

tokenizer_cls: tsr.models.tokenizers.triplane.Triplane1DTokenizer
tokenizer:
  plane_size: 32
  num_channels: 1024

backbone_cls: tsr.models.transformer.transformer_1d.Transformer1D
backbone:
  in_channels: ${tokenizer.num_channels}
  num_attention_heads: 16
  attention_head_dim: 64
  num_layers: 16
  cross_attention_dim: 768

post_processor_cls: tsr.models.network_utils.TriplaneUpsampleNetwork
post_processor:
  in_channels: 1024
  out_channels: 40

decoder_cls: tsr.models.network_utils.NeRFMLP
decoder:
  in_channels: 120 # 3 * 40
  n_neurons: 64
  n_hidden_layers: 9
  activation: silu

renderer_cls: tsr.models.nerf_renderer.TriplaneNeRFRenderer
renderer:
  radius: 0.87 # slightly larger than 0.5 * sqrt(3)
  feature_reduction: concat
  density_activation: exp
  density_bias: -1.0
  num_samples_per_ray: 128

Create mini dataset. Then try to fine tune TripoSR Model

import torch
import numpy as np
import os
from PIL import Image

import torch.optim as optim
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import transforms
from my_custom_dataset import CustomDataset  # Import your custom dataset class


def CustomModelTrain():  

    # Set up your training parameters
    device=['cpu','cuda'][torch.cuda.is_available()]
    new_checkpoint_file = './outs/ckpt/model0.ckpt'
    batch_size = 32
    learning_rate = 0.001
    num_epochs = 10
    epochs=10
    image_size=[512, 32][1]
    
     # Create your dataset and data loader
    # Replace 'transform' with your desired preprocessing transforms
    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    ])
  
    train_dataset = CustomDataset(transform=transform)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

    from tsr.system import TSR
    model = TSR.from_pretrained(
        "stabilityai/TripoSR",
        config_name="/system_path/TripoSR/outs/ckpt/config.yaml",
        weight_name="/system_path/TripoSR/outs/ckpt/model.ckpt",
        is_local=True,
    )
    
    # adjust the chunk size to balance between speed and memory usage
    model.renderer.set_chunk_size(8192)
    model.to(device)
    model.train()
        
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.0001)  
    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=list(range(0, epochs, epochs//10)), gamma=0.90)

    for epoch in range(0, epochs):
        
       #Please  Write traning Code 
       #-----------------------------------

    ## Try to Export New Checkpoint
    # checkpoint = {
    #     'model_state_dict': model.state_dict(),
    #     'optimizer_state_dict': optimizer.state_dict(),
    #     # Add any other information you want to save
    # }
    # torch.save(checkpoint, new_checkpoint_file)

if __name__ == '__main__': 
    CustomModelTrain()

Please help me @bennyguo @pookiefoof @kostarion @varunjampani to complete this script.

Thanks!