app.py

import gradio as gr
import os
import shutil
os.environ['SPCONV_ALGO'] = 'native'
from typing import *
import torch
import numpy as np
import imageio
import uuid
from easydict import EasyDict as edict
from PIL import Image
from trellis.pipelines import TrellisImageTo3DPipeline
from trellis.representations import Gaussian, MeshExtractResult
from trellis.utils import render_utils, postprocessing_utils
from gradio_litmodel3d import LitModel3D


def check_gpu():
    """Check if CUDA GPU is available and properly initialized"""
    if not torch.cuda.is_available():
        raise RuntimeError(
            "This application requires a CUDA-capable GPU to run. "
            "No CUDA GPU was detected in your system."
        )
    
    # Print GPU information for debugging
    gpu_count = torch.cuda.device_count()
    print(f"Found {gpu_count} CUDA GPU(s)")
    for i in range(gpu_count):
        gpu_name = torch.cuda.get_device_name(i)
        print(f"GPU {i}: {gpu_name}")

    # Try to initialize CUDA
    try:
        torch.cuda.init()
        current_device = torch.cuda.current_device()
        print(f"Using GPU {current_device}: {torch.cuda.get_device_name(current_device)}")
    except Exception as e:
        raise RuntimeError(f"Failed to initialize CUDA: {str(e)}")

# ... [rest of the code remains exactly the same until main] ...

if __name__ == "__main__":
    # Check GPU availability first
    check_gpu()
    
    # Initialize pipeline with explicit device setting
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    pipeline = TrellisImageTo3DPipeline.from_pretrained(
        "JeffreyXiang/TRELLIS-image-large"
    ).to(device)
    
    try:
        # Use smaller test image and explicit device
        test_img = np.zeros((256, 256, 3), dtype=np.uint8)
        pipeline.preprocess_image(Image.fromarray(test_img))
        del test_img
    except Exception as e:
        print(f"Warning: Failed to preload rembg: {str(e)}")
    
    # Launch the demo
    demo.launch()