# 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 )