Refine Gradio interface by simplifying API endpoint settings, maintaining minimal configurations for the Blocks interface and button actions to enhance usability and performance.
2fe5022
| # 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 | |
| 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) | |
| 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 | |
| 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:) | |
| ''' | |
| def generate(image, bg_choice, fg_ratio, bg_color, seed_val, guidance, steps): | |
| if image is None: | |
| raise gr.Error("No image uploaded!") | |
| processed = preprocess_image(image, bg_choice, fg_ratio, bg_color) | |
| return gen_image(processed, seed_val, guidance, steps) | |
| # Create a Blocks interface with minimal settings | |
| with gr.Blocks( | |
| analytics_enabled=False, | |
| title="CRM: Single Image to 3D Textured Mesh with Convolutional Reconstruction Model" | |
| ) as demo: | |
| gr.Markdown(_DESCRIPTION) | |
| with gr.Row(): | |
| with gr.Column(): | |
| input_image = gr.Image( | |
| label="Image input", | |
| image_mode="RGBA", | |
| sources="upload", | |
| type="pil", | |
| ) | |
| bg_choice = gr.Radio( | |
| ["Alpha as mask", "Auto Remove background"], | |
| value="Auto Remove background", | |
| label="background choice" | |
| ) | |
| fg_ratio = gr.Slider( | |
| label="Foreground Ratio", | |
| minimum=0.5, | |
| maximum=1.0, | |
| value=1.0, | |
| step=0.05, | |
| ) | |
| bg_color = gr.ColorPicker(label="Background Color", value="#7F7F7F", interactive=False) | |
| seed_val = gr.Number(value=1234, label="seed", precision=0) | |
| guidance = gr.Number(value=5.5, minimum=3, maximum=10, label="guidance_scale") | |
| steps = gr.Number(value=30, minimum=30, maximum=100, label="sample steps", precision=0) | |
| generate_btn = gr.Button("Generate") | |
| with gr.Column(): | |
| output_rgb = gr.Image(interactive=False, label="Output RGB image") | |
| output_ccm = gr.Image(interactive=False, label="Output CCM image") | |
| output_glb = gr.Model3D(label="Output GLB", interactive=False) | |
| # Connect the button click event | |
| generate_btn.click( | |
| fn=generate, | |
| inputs=[input_image, bg_choice, fg_ratio, bg_color, seed_val, guidance, steps], | |
| outputs=[output_rgb, output_ccm, output_glb], | |
| concurrency_limit=1 | |
| ) | |
| # Add examples without caching | |
| gr.Examples( | |
| examples=[[os.path.join("examples", i)] for i in os.listdir("examples")], | |
| inputs=input_image, | |
| cache_examples=False # Disable example caching | |
| ) | |
| # Launch with Spaces-specific settings | |
| if __name__ == "__main__": | |
| import os | |
| from spaces.zero.gradio import launch # Use Spaces specific launch function | |
| if os.environ.get("SPACE_ID") is not None: # We're running on Hugging Face Spaces | |
| launch( | |
| demo, | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| show_error=True, | |
| enable_queue=True, | |
| max_threads=1, | |
| api_name=None, # Disable API endpoint generation | |
| share=False, # Don't use share on Spaces | |
| prevent_thread_lock=True # Prevent thread lock issues | |
| ) | |
| else: # Local development | |
| demo.launch( | |
| server_name="0.0.0.0", | |
| server_port=7860, | |
| show_error=True, | |
| share=True, # Only use share=True for local development | |
| api_name=None # Disable API endpoint | |
| ) |