Spaces:
Sleeping
Sleeping
| import spaces | |
| import os | |
| import gradio as gr | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| import trimesh | |
| import random | |
| from transformers import AutoModelForImageSegmentation | |
| from torchvision import transforms | |
| from huggingface_hub import hf_hub_download, snapshot_download | |
| import subprocess | |
| import shutil | |
| # install others | |
| subprocess.run("pip install spandrel==0.4.1 --no-deps", shell=True, check=True) | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| DTYPE = torch.float16 | |
| print("DEVICE: ", DEVICE) | |
| DEFAULT_FACE_NUMBER = 100000 | |
| MAX_SEED = np.iinfo(np.int32).max | |
| TRIPOSG_REPO_URL = "https://github.com/VAST-AI-Research/TripoSG.git" | |
| MV_ADAPTER_REPO_URL = "https://github.com/huanngzh/MV-Adapter.git" | |
| RMBG_PRETRAINED_MODEL = "checkpoints/RMBG-1.4" | |
| TRIPOSG_PRETRAINED_MODEL = "checkpoints/TripoSG" | |
| TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tmp") | |
| os.makedirs(TMP_DIR, exist_ok=True) | |
| TRIPOSG_CODE_DIR = "./triposg" | |
| if not os.path.exists(TRIPOSG_CODE_DIR): | |
| os.system(f"git clone {TRIPOSG_REPO_URL} {TRIPOSG_CODE_DIR}") | |
| MV_ADAPTER_CODE_DIR = "./mv_adapter" | |
| if not os.path.exists(MV_ADAPTER_CODE_DIR): | |
| os.system(f"git clone {MV_ADAPTER_REPO_URL} {MV_ADAPTER_CODE_DIR} && cd {MV_ADAPTER_CODE_DIR} && git checkout 7d37a97e9bc223cdb8fd26a76bd8dd46504c7c3d") | |
| import sys | |
| sys.path.append(TRIPOSG_CODE_DIR) | |
| sys.path.append(os.path.join(TRIPOSG_CODE_DIR, "scripts")) | |
| sys.path.append(MV_ADAPTER_CODE_DIR) | |
| sys.path.append(os.path.join(MV_ADAPTER_CODE_DIR, "scripts")) | |
| HEADER = """ | |
| # 🎨 PolyGenixAI: Transform Ideas into 3D Masterpieces | |
| ## Unleash Your Creativity with AI-Powered 3D Generation by AnvilInteractive Solutions | |
| <p style="font-size: 1.1em;">By <a href="https://www.anvilinteractive.com/" style="color: #3B82F6; text-decoration: none; font-weight: bold;">AnvilInteractive Solutions</a></p> | |
| ## 🚀 Get Started: | |
| 1. **Upload an Image** (clear, single-object images work best) | |
| 2. **Select Filters** to customize styles | |
| 3. Click **Generate 3D Model** to create your mesh | |
| 4. Click **Apply Texture** to enhance with realistic textures | |
| 5. **Download GLB** to save your creation | |
| <p style="font-size: 0.9em; margin-top: 10px;">Powered by advanced AI and multi-view technology from AnvilInteractive Solutions. Join our community at <a href="https://www.anvilinteractive.com/polygenixai" style="color: #3B82F6; text-decoration: none;">PolyGenixAI</a> for tips and inspiration.</p> | |
| <style> | |
| .gr-button-primary { | |
| background-color: #3B82F6 !important; | |
| color: white !important; | |
| border-radius: 8px !important; | |
| padding: 10px 20px !important; | |
| font-weight: 600 !important; | |
| } | |
| .gr-button-secondary { | |
| background-color: #E5E7EB !important; | |
| color: #1F2937 !important; | |
| border-radius: 8px !important; | |
| padding: 10px 20px !important; | |
| } | |
| .gr-panel { | |
| border-radius: 12px !important; | |
| box-shadow: 0 4px 6px rgba(0,0,0,0.1) !important; | |
| padding: 20px !important; | |
| } | |
| .gr-accordion { | |
| background-color: #F9FAFB !important; | |
| border-radius: 8px !important; | |
| } | |
| </style> | |
| """ | |
| # triposg | |
| from image_process import prepare_image | |
| from briarmbg import BriaRMBG | |
| snapshot_download("briaai/RMBG-1.4", local_dir=RMBG_PRETRAINED_MODEL) | |
| rmbg_net = BriaRMBG.from_pretrained(RMBG_PRETRAINED_MODEL).to(DEVICE) | |
| rmbg_net.eval() | |
| from triposg.pipelines.pipeline_triposg import TripoSGPipeline | |
| snapshot_download("VAST-AI/TripoSG", local_dir=TRIPOSG_PRETRAINED_MODEL) | |
| triposg_pipe = TripoSGPipeline.from_pretrained(TRIPOSG_PRETRAINED_MODEL).to(DEVICE, DTYPE) | |
| # mv adapter | |
| NUM_VIEWS = 6 | |
| from inference_ig2mv_sdxl import prepare_pipeline, preprocess_image, remove_bg | |
| from mvadapter.utils import get_orthogonal_camera, tensor_to_image, make_image_grid | |
| from mvadapter.utils.render import NVDiffRastContextWrapper, load_mesh, render | |
| mv_adapter_pipe = prepare_pipeline( | |
| base_model="stabilityai/stable-diffusion-xl-base-1.0", | |
| vae_model="madebyollin/sdxl-vae-fp16-fix", | |
| unet_model=None, | |
| lora_model=None, | |
| adapter_path="huanngzh/mv-adapter", | |
| scheduler=None, | |
| num_views=NUM_VIEWS, | |
| device=DEVICE, | |
| dtype=torch.float16, | |
| ) | |
| birefnet = AutoModelForImageSegmentation.from_pretrained( | |
| "ZhengPeng7/BiRefNet", trust_remote_code=True | |
| ) | |
| birefnet.to(DEVICE) | |
| transform_image = transforms.Compose( | |
| [ | |
| transforms.Resize((1024, 1024)), | |
| transforms.ToTensor(), | |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), | |
| ] | |
| ) | |
| remove_bg_fn = lambda x: remove_bg(x, birefnet, transform_image, DEVICE) | |
| if not os.path.exists("checkpoints/RealESRGAN_x2plus.pth"): | |
| hf_hub_download("dtarnow/UPscaler", filename="RealESRGAN_x2plus.pth", local_dir="checkpoints") | |
| if not os.path.exists("checkpoints/big-lama.pt"): | |
| subprocess.run("wget -P checkpoints/ https://github.com/Sanster/models/releases/download/add_big_lama/big-lama.pt", shell=True, check=True) | |
| def start_session(req: gr.Request): | |
| save_dir = os.path.join(TMP_DIR, str(req.session_hash)) | |
| os.makedirs(save_dir, exist_ok=True) | |
| print("start session, mkdir", save_dir) | |
| def end_session(req: gr.Request): | |
| save_dir = os.path.join(TMP_DIR, str(req.session_hash)) | |
| shutil.rmtree(save_dir) | |
| def get_random_hex(): | |
| random_bytes = os.urandom(8) | |
| random_hex = random_bytes.hex() | |
| return random_hex | |
| def get_random_seed(randomize_seed, seed): | |
| if randomize_seed: | |
| seed = random.randint(0, MAX_SEED) | |
| return seed | |
| def run_full(image: str, req: gr.Request): | |
| seed = 0 | |
| num_inference_steps = 50 | |
| guidance_scale = 7.5 | |
| simplify = True | |
| target_face_num = DEFAULT_FACE_NUMBER | |
| image_seg = prepare_image(image, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net) | |
| outputs = triposg_pipe( | |
| image=image_seg, | |
| generator=torch.Generator(device=triposg_pipe.device).manual_seed(seed), | |
| num_inference_steps=num_inference_steps, | |
| guidance_scale=guidance_scale | |
| ).samples[0] | |
| print("mesh extraction done") | |
| mesh = trimesh.Trimesh(outputs[0].astype(np.float32), np.ascontiguousarray(outputs[1])) | |
| if simplify: | |
| print("start simplify") | |
| from utils import simplify_mesh | |
| mesh = simplify_mesh(mesh, target_face_num) | |
| save_dir = os.path.join(TMP_DIR, "examples") | |
| os.makedirs(save_dir, exist_ok=True) | |
| mesh_path = os.path.join(save_dir, f"polygenixai_{get_random_hex()}.glb") | |
| mesh.export(mesh_path) | |
| print("save to ", mesh_path) | |
| torch.cuda.empty_cache() | |
| height, width = 768, 768 | |
| # Prepare cameras | |
| cameras = get_orthogonal_camera( | |
| elevation_deg=[0, 0, 0, 0, 89.99, -89.99], | |
| distance=[1.8] * NUM_VIEWS, | |
| left=-0.55, | |
| right=0.55, | |
| bottom=-0.55, | |
| top=0.55, | |
| azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]], | |
| device=DEVICE, | |
| ) | |
| ctx = NVDiffRastContextWrapper(device=DEVICE, context_type="cuda") | |
| mesh = load_mesh(mesh_path, rescale=True, device=DEVICE) | |
| render_out = render( | |
| ctx, | |
| mesh, | |
| cameras, | |
| height=height, | |
| width=width, | |
| render_attr=False, | |
| normal_background=0.0, | |
| ) | |
| control_images = ( | |
| torch.cat( | |
| [ | |
| (render_out.pos + 0.5).clamp(0, 1), | |
| (render_out.normal / 2 + 0.5).clamp(0, 1), | |
| ], | |
| dim=-1, | |
| ) | |
| .permute(0, 3, 1, 2) | |
| .to(DEVICE) | |
| ) | |
| image = Image.open(image) | |
| image = remove_bg_fn(image) | |
| image = preprocess_image(image, height, width) | |
| pipe_kwargs = {} | |
| if seed != -1 and isinstance(seed, int): | |
| pipe_kwargs["generator"] = torch.Generator(device=DEVICE).manual_seed(seed) | |
| images = mv_adapter_pipe( | |
| "high quality", | |
| height=height, | |
| width=width, | |
| num_inference_steps=15, | |
| guidance_scale=3.0, | |
| num_images_per_prompt=NUM_VIEWS, | |
| control_image=control_images, | |
| control_conditioning_scale=1.0, | |
| reference_image=image, | |
| reference_conditioning_scale=1.0, | |
| negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast", | |
| cross_attention_kwargs={"scale": 1.0}, | |
| **pipe_kwargs, | |
| ).images | |
| torch.cuda.empty_cache() | |
| mv_image_path = os.path.join(save_dir, f"polygenixai_mv_{get_random_hex()}.png") | |
| make_image_grid(images, rows=1).save(mv_image_path) | |
| from texture import TexturePipeline, ModProcessConfig | |
| texture_pipe = TexturePipeline( | |
| upscaler_ckpt_path="checkpoints/RealESRGAN_x2plus.pth", | |
| inpaint_ckpt_path="checkpoints/big-lama.pt", | |
| device=DEVICE, | |
| ) | |
| textured_glb_path = texture_pipe( | |
| mesh_path=mesh_path, | |
| save_dir=save_dir, | |
| save_name=f"polygenixai_texture_mesh_{get_random_hex()}.glb", | |
| uv_unwarp=True, | |
| uv_size=4096, | |
| rgb_path=mv_image_path, | |
| rgb_process_config=ModProcessConfig(view_upscale=True, inpaint_mode="view"), | |
| camera_azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]], | |
| ) | |
| return image_seg, mesh_path, textured_glb_path | |
| def run_segmentation(image: str): | |
| image = prepare_image(image, bg_color=np.array([1.0, 1.0, 1.0]), rmbg_net=rmbg_net) | |
| return image | |
| def image_to_3d( | |
| image: Image.Image, | |
| seed: int, | |
| num_inference_steps: int, | |
| guidance_scale: float, | |
| simplify: bool, | |
| target_face_num: int, | |
| req: gr.Request | |
| ): | |
| outputs = triposg_pipe( | |
| image=image, | |
| generator=torch.Generator(device=triposg_pipe.device).manual_seed(seed), | |
| num_inference_steps=num_inference_steps, | |
| guidance_scale=guidance_scale | |
| ).samples[0] | |
| print("mesh extraction done") | |
| mesh = trimesh.Trimesh(outputs[0].astype(np.float32), np.ascontiguousarray(outputs[1])) | |
| if simplify: | |
| print("start simplify") | |
| from utils import simplify_mesh | |
| mesh = simplify_mesh(mesh, target_face_num) | |
| save_dir = os.path.join(TMP_DIR, str(req.session_hash)) | |
| mesh_path = os.path.join(save_dir, f"polygenixai_{get_random_hex()}.glb") | |
| mesh.export(mesh_path) | |
| print("save to ", mesh_path) | |
| torch.cuda.empty_cache() | |
| return mesh_path | |
| def run_texture(image: Image, mesh_path: str, seed: int, req: gr.Request): | |
| height, width = 768, 768 | |
| # Prepare cameras | |
| cameras = get_orthogonal_camera( | |
| elevation_deg=[0, 0, 0, 0, 89.99, -89.99], | |
| distance=[1.8] * NUM_VIEWS, | |
| left=-0.55, | |
| right=0.55, | |
| bottom=-0.55, | |
| top=0.55, | |
| azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]], | |
| device=DEVICE, | |
| ) | |
| ctx = NVDiffRastContextWrapper(device=DEVICE, context_type="cuda") | |
| mesh = load_mesh(mesh_path, rescale=True, device=DEVICE) | |
| render_out = render( | |
| ctx, | |
| mesh, | |
| cameras, | |
| height=height, | |
| width=width, | |
| render_attr=False, | |
| normal_background=0.0, | |
| ) | |
| control_images = ( | |
| torch.cat( | |
| [ | |
| (render_out.pos + 0.5).clamp(0, 1), | |
| (render_out.normal / 2 + 0.5).clamp(0, 1), | |
| ], | |
| dim=-1, | |
| ) | |
| .permute(0, 3, 1, 2) | |
| .to(DEVICE) | |
| ) | |
| image = Image.open(image) | |
| image = remove_bg_fn(image) | |
| image = preprocess_image(image, height, width) | |
| pipe_kwargs = {} | |
| if seed != -1 and isinstance(seed, int): | |
| pipe_kwargs["generator"] = torch.Generator(device=DEVICE).manual_seed(seed) | |
| images = mv_adapter_pipe( | |
| "high quality", | |
| height=height, | |
| width=width, | |
| num_inference_steps=15, | |
| guidance_scale=3.0, | |
| num_images_per_prompt=NUM_VIEWS, | |
| control_image=control_images, | |
| control_conditioning_scale=1.0, | |
| reference_image=image, | |
| reference_conditioning_scale=1.0, | |
| negative_prompt="watermark, ugly, deformed, noisy, blurry, low contrast", | |
| cross_attention_kwargs={"scale": 1.0}, | |
| **pipe_kwargs, | |
| ).images | |
| torch.cuda.empty_cache() | |
| save_dir = os.path.join(TMP_DIR, str(req.session_hash)) | |
| mv_image_path = os.path.join(save_dir, f"polygenixai_mv_{get_random_hex()}.png") | |
| make_image_grid(images, rows=1).save(mv_image_path) | |
| from texture import TexturePipeline, ModProcessConfig | |
| texture_pipe = TexturePipeline( | |
| upscaler_ckpt_path="checkpoints/RealESRGAN_x2plus.pth", | |
| inpaint_ckpt_path="checkpoints/big-lama.pt", | |
| device=DEVICE, | |
| ) | |
| textured_glb_path = texture_pipe( | |
| mesh_path=mesh_path, | |
| save_dir=save_dir, | |
| save_name=f"polygenixai_texture_mesh_{get_random_hex()}.glb", | |
| uv_unwarp=True, | |
| uv_size=4096, | |
| rgb_path=mv_image_path, | |
| rgb_process_config=ModProcessConfig(view_upscale=True, inpaint_mode="view"), | |
| camera_azimuth_deg=[x - 90 for x in [0, 90, 180, 270, 180, 180]], | |
| ) | |
| return textured_glb_path | |
| with gr.Blocks(title="PolyGenixAI", css="body { background-color: #F3F4F6; } .gr-panel { background-color: white; }") as demo: | |
| gr.Markdown(HEADER) | |
| with gr.Tabs(): | |
| with gr.Tab("Create 3D Model"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| image_prompts = gr.Image(label="Upload Image", type="filepath", height=300) | |
| seg_image = gr.Image(label="Preview Segmentation", type="pil", format="png", interactive=False, height=300) | |
| with gr.Accordion("Style & Settings", open=True): | |
| style_filter = gr.Dropdown( | |
| choices=["None", "Realistic", "Fantasy", "Cartoon", "Sci-Fi", "Vintage"], | |
| label="Style Filter", | |
| value="None", | |
| info="Select a style to enhance your 3D model (optional)" | |
| ) | |
| seed = gr.Slider( | |
| label="Seed", | |
| minimum=0, | |
| maximum=MAX_SEED, | |
| step=1, | |
| value=0 | |
| ) | |
| randomize_seed = gr.Checkbox(label="Randomize Seed", value=True) | |
| num_inference_steps = gr.Slider( | |
| label="Inference Steps", | |
| minimum=8, | |
| maximum=50, | |
| step=1, | |
| value=50, | |
| info="Higher steps improve quality but take longer" | |
| ) | |
| guidance_scale = gr.Slider( | |
| label="Guidance Scale", | |
| minimum=0.0, | |
| maximum=20.0, | |
| step=0.1, | |
| value=7.0, | |
| info="Controls how closely the model follows the input" | |
| ) | |
| reduce_face = gr.Checkbox(label="Simplify Mesh", value=True) | |
| target_face_num = gr.Slider( | |
| maximum=1000000, | |
| minimum=10000, | |
| value=DEFAULT_FACE_NUMBER, | |
| label="Target Face Number", | |
| info="Adjust mesh complexity" | |
| ) | |
| gen_button = gr.Button("Generate 3D Model", variant="primary") | |
| gen_texture_button = gr.Button("Apply Texture", variant="secondary", interactive=False) | |
| with gr.Column(scale=1): | |
| model_output = gr.Model3D(label="3D Model Preview", interactive=False, height=400) | |
| textured_model_output = gr.Model3D(label="Textured 3D Model", interactive=False, height=400) | |
| download_button = gr.Button("Download GLB", variant="secondary") | |
| with gr.Tab("Gallery & Community"): | |
| gr.Markdown("### Explore Creations") | |
| gr.Examples( | |
| examples=[ | |
| f"{TRIPOSG_CODE_DIR}/assets/example_data/{image}" | |
| for image in os.listdir(f"{TRIPOSG_CODE_DIR}/assets/example_data") | |
| ], | |
| fn=run_full, | |
| inputs=[image_prompts], | |
| outputs=[seg_image, model_output, textured_model_output], | |
| cache_examples=True, | |
| ) | |
| gr.Markdown("Join our [PolyGenixAI Community](https://www.anvilinteractive.com/community) to share your creations and get inspired!") | |
| gen_button.click( | |
| run_segmentation, | |
| inputs=[image_prompts], | |
| outputs=[seg_image] | |
| ).then( | |
| get_random_seed, | |
| inputs=[randomize_seed, seed], | |
| outputs=[seed], | |
| ).then( | |
| image_to_3d, | |
| inputs=[ | |
| seg_image, | |
| seed, | |
| num_inference_steps, | |
| guidance_scale, | |
| reduce_face, | |
| target_face_num | |
| ], | |
| outputs=[model_output] | |
| ).then(lambda: gr.Button(interactive=True), outputs=[gen_texture_button]) | |
| gen_texture_button.click( | |
| run_texture, | |
| inputs=[image_prompts, model_output, seed], | |
| outputs=[textured_model_output] | |
| ) | |
| demo.load(start_session) | |
| demo.unload(end_session) | |
| demo.launch() |