app.py

# Fixed version with proper error handling and compatibility
import spaces
import argparse
import numpy as np
import gradio as gr
from omegaconf import OmegaConf
import torch
from PIL import Image
import PIL
from pipelines import TwoStagePipeline
from huggingface_hub import hf_hub_download
import os
import rembg
from typing import Any
import json
import os
import json
import argparse

from model import CRM
from inference import generate3d

# Move model initialization into a function that will be called by workers
def init_model():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--stage1_config",
        type=str,
        default="configs/nf7_v3_SNR_rd_size_stroke.yaml",
        help="config for stage1",
    )
    parser.add_argument(
        "--stage2_config",
        type=str,
        default="configs/stage2-v2-snr.yaml",
        help="config for stage2",
    )
    parser.add_argument("--device", type=str, default="cuda")
    args = parser.parse_args(args=[])  # Fix: provide empty args list

    # Download model files
    crm_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="CRM.pth")
    specs = json.load(open("configs/specs_objaverse_total.json"))
    model = CRM(specs)
    model.load_state_dict(torch.load(crm_path, map_location="cpu"), strict=False)
    model = model.to(args.device)

    # Load configs
    stage1_config = OmegaConf.load(args.stage1_config).config
    stage2_config = OmegaConf.load(args.stage2_config).config
    stage2_sampler_config = stage2_config.sampler
    stage1_sampler_config = stage1_config.sampler

    stage1_model_config = stage1_config.models
    stage2_model_config = stage2_config.models

    xyz_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="ccm-diffusion.pth")
    pixel_path = hf_hub_download(repo_id="Zhengyi/CRM", filename="pixel-diffusion.pth")
    stage1_model_config.resume = pixel_path
    stage2_model_config.resume = xyz_path

    pipeline = TwoStagePipeline(
        stage1_model_config,
        stage2_model_config,
        stage1_sampler_config,
        stage2_sampler_config,
        device=args.device,
        dtype=torch.float32
    )
    
    return model, pipeline, args

# Global variables to store model and pipeline
model = None
pipeline = None
args = None

@spaces.GPU
def get_model():
    """Lazy initialization of model and pipeline"""
    global model, pipeline, args
    if model is None or pipeline is None:
        model, pipeline, args = init_model()
    return model, pipeline

rembg_session = rembg.new_session()

def expand_to_square(image, bg_color=(0, 0, 0, 0)):
    # expand image to 1:1
    width, height = image.size
    if width == height:
        return image
    new_size = (max(width, height), max(width, height))
    new_image = Image.new("RGBA", new_size, bg_color)
    paste_position = ((new_size[0] - width) // 2, (new_size[1] - height) // 2)
    new_image.paste(image, paste_position)
    return new_image

def check_input_image(input_image):
    """Check if the input image is valid"""
    if input_image is None:
        raise gr.Error("No image uploaded!")
    return input_image

def remove_background(
    image: PIL.Image.Image,
    rembg_session: Any = None,
    force: bool = False,
    **rembg_kwargs,
) -> PIL.Image.Image:
    do_remove = True
    if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
        # explain why current do not rm bg
        print("alpha channel not empty, skip remove background, using alpha channel as mask")
        background = Image.new("RGBA", image.size, (0, 0, 0, 0))
        image = Image.alpha_composite(background, image)
        do_remove = False
    do_remove = do_remove or force
    if do_remove:
        image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
    return image

def do_resize_content(original_image: Image, scale_rate):
    # resize image content while retaining the original image size
    if scale_rate != 1:
        # Calculate the new size after rescaling
        new_size = tuple(int(dim * scale_rate) for dim in original_image.size)
        # Resize the image while maintaining the aspect ratio
        resized_image = original_image.resize(new_size)
        # Create a new image with the original size and black background
        padded_image = Image.new("RGBA", original_image.size, (0, 0, 0, 0))
        paste_position = ((original_image.width - resized_image.width) // 2, (original_image.height - resized_image.height) // 2)
        padded_image.paste(resized_image, paste_position)
        return padded_image
    else:
        return original_image

def add_background(image, bg_color=(255, 255, 255)):
    # given an RGBA image, alpha channel is used as mask to add background color
    background = Image.new("RGBA", image.size, bg_color)
    return Image.alpha_composite(background, image)

def add_random_background(image, color):
    # Add a random background to the image
    width, height = image.size
    background = Image.new("RGBA", image.size, color)
    return Image.alpha_composite(background, image)

@spaces.GPU
def preprocess_image(input_image, background_choice, foreground_ratio, back_ground_color):
    """Preprocess the input image"""
    try:
        # Check if image is provided
        if input_image is None:
            raise gr.Error("No image uploaded!")
            
        # Convert to PIL Image if needed
        if isinstance(input_image, np.ndarray):
            input_image = Image.fromarray(input_image)
        
        # Ensure RGBA mode
        if input_image.mode != "RGBA":
            input_image = input_image.convert("RGBA")
        
        # Process background
        if background_choice == "Auto Remove background":
            input_image = remove_background(input_image, rembg_session)
        elif background_choice == "Custom Background":
            input_image = add_random_background(input_image, back_ground_color)
            
        # Resize content if needed
        if foreground_ratio != 1.0:
            input_image = do_resize_content(input_image, foreground_ratio)
            
        return input_image
    except Exception as e:
        print(f"Error in preprocess_image: {str(e)}")
        raise gr.Error(f"Preprocessing failed: {str(e)}")

@spaces.GPU
def gen_image(processed_image, seed, scale, step):
    """Generate the 3D model"""
    try:
        # Get model and pipeline when needed
        model, pipeline = get_model()
        
        # Check if image is provided
        if processed_image is None:
            raise gr.Error("No processed image provided!")
        
        # Convert to numpy array
        if isinstance(processed_image, Image.Image):
            np_image = np.array(processed_image)
        else:
            np_image = processed_image
        
        # Set random seed
        torch.manual_seed(int(seed))
        np.random.seed(int(seed))
        
        # Generate images
        np_imgs, np_xyzs = pipeline.generate(
            np_image,
            guidance_scale=float(scale),
            num_inference_steps=int(step)
        )
        
        # Generate 3D model
        glb_path = generate3d(model, np_imgs, np_xyzs, args.device)
        
        return Image.fromarray(np_imgs), Image.fromarray(np_xyzs), glb_path
    except Exception as e:
        print(f"Error in gen_image: {str(e)}")
        raise gr.Error(f"Generation failed: {str(e)}")

def process_and_generate(image, bg_choice, fg_ratio, bg_color, seed_val, guidance, steps):
    """Combined function to process image and generate 3D model"""
    try:
        if image is None:
            raise gr.Error("No image uploaded!")
        
        # Preprocess the image
        processed = preprocess_image(image, bg_choice, fg_ratio, bg_color)
        
        # Generate 3D model
        rgb_img, ccm_img, glb_file = gen_image(processed, seed_val, guidance, steps)
        
        return processed, rgb_img, ccm_img, glb_file
    except Exception as e:
        print(f"Error in process_and_generate: {str(e)}")
        raise gr.Error(f"Process failed: {str(e)}")

_DESCRIPTION = '''
* Our [official implementation](https://github.com/thu-ml/CRM) uses UV texture instead of vertex color. It has better texture than this online demo.
* Project page of CRM: https://ml.cs.tsinghua.edu.cn/~zhengyi/CRM/
* If you find the output unsatisfying, try using different seeds:)
'''

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# CRM: Single Image to 3D Textured Mesh with Convolutional Reconstruction Model")
    gr.Markdown(_DESCRIPTION)
    with gr.Row():
        with gr.Column():
            with gr.Row():
                image_input = gr.Image(
                    label="Image input",
                    image_mode="RGBA",
                    sources="upload",
                    type="pil",
                )
                processed_image = gr.Image(label="Processed Image", interactive=False, type="pil", image_mode="RGB")
            with gr.Row():
                with gr.Column():
                    with gr.Row():
                        background_choice = gr.Radio([
                            "Alpha as mask",
                            "Auto Remove background"
                        ], value="Auto Remove background",
                        label="Background choice")
                    back_ground_color = gr.ColorPicker(label="Background Color", value="#7F7F7F", interactive=False)
                    foreground_ratio = gr.Slider(
                        label="Foreground Ratio",
                        minimum=0.5,
                        maximum=1.0,
                        value=1.0,
                        step=0.05,
                    )

                with gr.Column():
                    seed = gr.Number(value=1234, label="Seed", precision=0)
                    guidance_scale = gr.Number(value=5.5, minimum=3, maximum=10, label="Guidance scale")
                    step = gr.Number(value=30, minimum=30, maximum=100, label="Sample steps", precision=0)
            text_button = gr.Button("Generate 3D shape")
            if os.path.exists("examples") and os.listdir("examples"):
                gr.Examples(
                    examples=[os.path.join("examples", i) for i in os.listdir("examples") if i.lower().endswith(('.png', '.jpg', '.jpeg'))],
                    inputs=[image_input],
                    examples_per_page=20,
                )
        with gr.Column():
            image_output = gr.Image(interactive=False, label="Output RGB image")
            xyz_output = gr.Image(interactive=False, label="Output CCM image")
            output_model = gr.Model3D(
                label="Output GLB",
                interactive=False,
            )
            gr.Markdown("Note: Ensure that the input image is correctly pre-processed into a grey background, otherwise the results will be unpredictable.")

    inputs = [
        processed_image,
        seed,
        guidance_scale,
        step,
    ]
    outputs = [
        image_output,
        xyz_output,
        output_model,
    ]

    text_button.click(fn=check_input_image, inputs=[image_input]).success(
        fn=preprocess_image,
        inputs=[image_input, background_choice, foreground_ratio, back_ground_color],
        outputs=[processed_image],
    ).success(
        fn=gen_image,
        inputs=inputs,
        outputs=outputs,
    )

if __name__ == "__main__":
    demo.queue().launch()