app.py

# Not ready to use yet
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()

    # 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

@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("alhpa channl not enpty, 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 wile retain 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_groud_color):
    """Preprocess the input image"""
    try:
        # Get model and pipeline when needed
        model, pipeline = get_model()
        
        # Convert to numpy array
        np_image = np.array(input_image)
        
        # Process background
        if background_choice == "Remove Background":
            np_image = rembg.remove(np_image, session=rembg_session)
        elif background_choice == "Custom Background":
            np_image = add_random_background(np_image, back_groud_color)
            
        # Resize content if needed
        if foreground_ratio != 1.0:
            np_image = do_resize_content(Image.fromarray(np_image), foreground_ratio)
            np_image = np.array(np_image)
            
        return Image.fromarray(np_image)
    except Exception as e:
        print(f"Error in preprocess_image: {str(e)}")
        raise 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()
        
        # Convert to numpy array
        np_image = np.array(processed_image)
        
        # Set random seed
        torch.manual_seed(seed)
        np.random.seed(seed)
        
        # Generate images
        np_imgs, np_xyzs = pipeline.generate(
            np_image,
            guidance_scale=scale,
            num_inference_steps=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 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 Gradio interface
with gr.Blocks(title="CRM: 3D Character Generation from Single Image") as demo:
    gr.Markdown(_DESCRIPTION)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input Image", type="pil")
            background_choice = gr.Radio(
                choices=["Remove Background", "Custom Background"],
                value="Remove Background",
                label="Background Option"
            )
            foreground_ratio = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=1.0,
                step=0.1,
                label="Foreground Ratio"
            )
            back_groud_color = gr.ColorPicker(
                label="Background Color",
                value="#FFFFFF"
            )
            seed = gr.Number(
                label="Seed",
                value=42,
                precision=0
            )
            scale = gr.Slider(
                minimum=1.0,
                maximum=20.0,
                value=7.5,
                step=0.1,
                label="Guidance Scale"
            )
            step = gr.Slider(
                minimum=1,
                maximum=100,
                value=50,
                step=1,
                label="Steps"
            )
            generate_btn = gr.Button("Generate 3D Model")
            
        with gr.Column():
            processed_image = gr.Image(label="Processed Image", type="pil")
            output_image = gr.Image(label="Generated Image", type="pil")
            output_xyz = gr.Image(label="Generated XYZ", type="pil")
            output_glb = gr.Model3D(label="Generated 3D Model")
    
    # Connect the functions with explicit API names
    generate_btn.click(
        fn=check_input_image,
        inputs=[input_image],
        outputs=[input_image],
        api_name="check_input_image"
    ).success(
        fn=preprocess_image,
        inputs=[input_image, background_choice, foreground_ratio, back_groud_color],
        outputs=[processed_image],
        api_name="preprocess_image"
    ).success(
        fn=gen_image,
        inputs=[processed_image, seed, scale, step],
        outputs=[output_image, output_xyz, output_glb],
        api_name="gen_image"
    )

# For Hugging Face Spaces, use minimal configuration
demo.queue().launch(
    show_error=True  # Only keep error display for debugging
)