import gradio as gr import torch import numpy as np from functools import partial from typing import Optional from shap_e.diffusion.gaussian_diffusion import diffusion_from_config from shap_e.diffusion.sample import sample_latents from shap_e.models.download import load_model, load_config from shap_e.util.notebooks import create_pan_cameras, decode_latent_mesh import trimesh import torch.nn as nn import os import random import warnings from huggingface_hub import hf_hub_download import hashlib import sys sys.tracebacklimit = 0 def set_seed(seed=1024): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) torch.backends.cudnn.deterministic = True def freeze_params(params): for param in params: param.requires_grad = False class Blocks(gr.Blocks): def __init__( self, theme: str = "default", analytics_enabled: Optional[bool] = None, mode: str = "blocks", title: str = "Gradio", css: Optional[str] = None, **kwargs, ): self.extra_configs = { 'thumbnail': kwargs.pop('thumbnail', ''), 'url': kwargs.pop('url', 'https://gradio.app/'), 'creator': kwargs.pop('creator', '@teamGradio'), } super(Blocks, self).__init__(theme, analytics_enabled, mode, title, css, **kwargs) warnings.filterwarnings("ignore") def get_config_file(self): config = super(Blocks, self).get_config_file() for k, v in self.extra_configs.items(): config[k] = v return config def optimize_all(xm, models, initial_noise, noise_start_t, diffusion, latent_model, device, prompt, instruction, rand_seed): state = {} out_gen_1, out_gen_2, out_gen_3, out_gen_4, state = generate_3d_with_shap_e(xm, diffusion, latent_model, device, prompt, rand_seed, state) edited_1, edited_2, edited_3, edited_4, state = _3d_editing(xm, models, diffusion, initial_noise, noise_start_t, device, instruction, rand_seed, state) print(state) return out_gen_1, out_gen_2, out_gen_3, out_gen_4, edited_1, edited_2, edited_3, edited_4 def generate_3d_with_shap_e(xm, diffusion, latent_model, device, prompt, rand_seed, state): set_seed(rand_seed) batch_size = 4 guidance_scale = 15.0 xm.renderer.volume.bbox_max = torch.tensor([1.0, 1.0, 1.0]).to(device) xm.renderer.volume.bbox_min = torch.tensor([-1.0, -1.0, -1.0]).to(device) xm.renderer.volume.bbox = torch.stack([xm.renderer.volume.bbox_min, xm.renderer.volume.bbox_max]) print("prompt: ", prompt, "rand_seed: ", rand_seed, "state:", state) latents = sample_latents( batch_size=batch_size, model=latent_model, diffusion=diffusion, guidance_scale=guidance_scale, model_kwargs=dict(texts=[prompt] * batch_size), progress=True, clip_denoised=True, use_fp16=True, use_karras=True, karras_steps=64, sigma_min=1e-3, sigma_max=160, s_churn=0, ) prompt_hash = str(hashlib.sha256((prompt + '_' + str(rand_seed)).encode('utf-8')).hexdigest()) mesh_path = [] output_path = './logs' os.makedirs(os.path.join(output_path, 'source'), exist_ok=True) state['latent'] = [] state['prompt'] = prompt state['rand_seed_1'] = rand_seed for i, latent in enumerate(latents): output_path_tmp = os.path.join(output_path, 'source', '{}_{}.obj'.format(prompt_hash, i)) t_obj = decode_latent_mesh(xm, latent).tri_mesh() with open(output_path_tmp, 'w') as f: t_obj.write_obj(f) mesh = trimesh.load_mesh(output_path_tmp) angle = np.radians(180) axis = [0, 1, 0] rotation_matrix = trimesh.transformations.rotation_matrix(angle, axis) mesh.apply_transform(rotation_matrix) angle = np.radians(90) axis = [1, 0, 0] rotation_matrix = trimesh.transformations.rotation_matrix(angle, axis) mesh.apply_transform(rotation_matrix) output_path_tmp = os.path.join(output_path, 'source', '{}_{}.obj'.format(prompt_hash, i)) mesh.export(output_path_tmp) state['latent'].append(latent.clone().detach()) mesh_path.append(output_path_tmp) return mesh_path[0], mesh_path[1], mesh_path[2], mesh_path[3], state def _3d_editing(xm, models, diffusion, initial_noise, start_t, device, instruction, rand_seed, state): set_seed(rand_seed) mesh_path = [] prompt = state['prompt'] rand_seed_1 = state['rand_seed_1'] print("prompt: ", prompt, "rand_seed: ", rand_seed, "instruction:", instruction, "state:", state) prompt_hash = str(hashlib.sha256((prompt + '_' + str(rand_seed_1) + '_' + instruction + '_' + str(rand_seed)).encode('utf-8')).hexdigest()) if 'santa' in instruction: e_type = 'santa_hat' elif 'rainbow' in instruction: e_type = 'rainbow' elif 'gold' in instruction: e_type = 'golden' elif 'lego' in instruction: e_type = 'lego' elif 'wooden' in instruction: e_type = 'wooden' elif 'cyber' in instruction: e_type = 'cyber' # import pdb; pdb.set_trace() model = models[e_type].to(device) noise_initial = initial_noise[e_type].to(device) noise_start_t = start_t[e_type] general_save_path = './logs/edited' os.makedirs(general_save_path, exist_ok=True) for i, latent in enumerate(state['latent']): latent = latent.to(device) text_embeddings_clip = model.cached_model_kwargs(1, dict(texts=[instruction])) print("shape of latent: ", latent.clone().unsqueeze(0).shape, "instruction: ", instruction) ref_latent = latent.clone().unsqueeze(0) t_1 = torch.randint(noise_start_t, noise_start_t + 1, (1,), device=device).long() noise_input = diffusion.q_sample(ref_latent, t_1, noise=noise_initial) out_1 = diffusion.p_mean_variance(model, noise_input, t_1, clip_denoised=True, model_kwargs=text_embeddings_clip, condition_latents=ref_latent) updated_latents = out_1['pred_xstart'] if 'santa' in instruction: xm.renderer.volume.bbox_max = torch.tensor([1.0, 1.0, 1.25]).to(device) xm.renderer.volume.bbox_min = torch.tensor([-1.0, -1.0, -1]).to(device) xm.renderer.volume.bbox = torch.stack([xm.renderer.volume.bbox_min, xm.renderer.volume.bbox_max]) else: xm.renderer.volume.bbox_max = torch.tensor([1.0, 1.0, 1.0]).to(device) xm.renderer.volume.bbox_min = torch.tensor([-1.0, -1.0, -1.0]).to(device) xm.renderer.volume.bbox = torch.stack([xm.renderer.volume.bbox_min, xm.renderer.volume.bbox_max]) for latent_idx, updated_latent in enumerate(updated_latents): output_path = os.path.join(general_save_path, '{}_{}.obj'.format(prompt_hash, i)) t = decode_latent_mesh(xm, updated_latent).tri_mesh() with open(output_path, 'w') as f: t.write_obj(f) mesh = trimesh.load_mesh(output_path) angle = np.radians(180) axis = [0, 1, 0] rotation_matrix = trimesh.transformations.rotation_matrix(angle, axis) mesh.apply_transform(rotation_matrix) angle = np.radians(90) axis = [1, 0, 0] rotation_matrix = trimesh.transformations.rotation_matrix(angle, axis) mesh.apply_transform(rotation_matrix) output_path = os.path.join(general_save_path, '{}_{}.obj'.format(prompt_hash, i)) mesh.export(output_path) mesh_path.append(output_path) return mesh_path[0], mesh_path[1], mesh_path[2], mesh_path[3], state def main(): css = """ #img2img_image, #img2img_image > .fixed-height, #img2img_image > .fixed-height > div, #img2img_image > .fixed-height > div > img { height: var(--height) !important; max-height: var(--height) !important; min-height: var(--height) !important; } #paper-info a { color:#008AD7; text-decoration: none; } #paper-info a:hover { cursor: pointer; text-decoration: none; } .tooltip { color: #555; position: relative; display: inline-block; cursor: pointer; } .tooltip .tooltiptext { visibility: hidden; width: 400px; background-color: #555; color: #fff; text-align: center; padding: 5px; border-radius: 5px; position: absolute; z-index: 1; /* Set z-index to 1 */ left: 10px; top: 100%; opacity: 0; transition: opacity 0.3s; } .tooltip:hover .tooltiptext { visibility: visible; opacity: 1; z-index: 9999; /* Set a high z-index value when hovering */ } """ rescale_js = """ function(x) { const root = document.querySelector('gradio-app').shadowRoot || document.querySelector('gradio-app'); let image_scale = parseFloat(root.querySelector('#image_scale input').value) || 1.0; const image_width = root.querySelector('#img2img_image').clientWidth; const target_height = parseInt(image_width * image_scale); document.body.style.setProperty('--height', `${target_height}px`); root.querySelectorAll('button.justify-center.rounded')[0].style.display='none'; root.querySelectorAll('button.justify-center.rounded')[1].style.display='none'; return x; } """ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') latent_model = load_model('text300M', device=device) xm = load_model('transmitter', device=device) diffusion = diffusion_from_config(load_config('diffusion')) freeze_params(xm.parameters()) models = dict() initial_noise = dict() noise_start_t = dict() editing_types = ['rainbow', 'santa_hat', 'lego', 'golden', 'wooden', 'cyber'] for editing_type in editing_types: tmp_model = load_model('text300M', device=device) with torch.no_grad(): new_proj = nn.Linear(1024 * 2, 1024, device=device, dtype=tmp_model.wrapped.input_proj.weight.dtype) new_proj.weight = nn.Parameter(torch.zeros_like(new_proj.weight)) new_proj.weight[:, :1024].copy_(tmp_model.wrapped.input_proj.weight) # new_proj.bias = nn.Parameter(torch.zeros_like(new_proj.bias)) new_proj.bias[:1024].copy_(tmp_model.wrapped.input_proj.bias) tmp_model.wrapped.input_proj = new_proj ckp = torch.load(hf_hub_download(repo_id='silentchen/Shap_Editor', subfolder='single', filename='{}.pt'.format(editing_type)), map_location='cpu') tmp_model.load_state_dict(ckp['model']) noise_initial = ckp['initial_noise']['noise'].to(device) initial_noise[editing_type] = noise_initial noise_start_t[editing_type] = ckp['t_start'] models[editing_type] = tmp_model with Blocks( css=css, analytics_enabled=False, title="SHAPE-EDITOR demo", ) as demo: description = """

SHAP-EDITOR: Instruction-guided
Latent 3D Editing in Seconds
[Project Page] [Paper] [GitHub]

""" state = gr.State({}) gr.HTML(description) with gr.Column(): with gr.Column(): gr.HTML('Step 1: generate original 3D object using Shap-E.') prompt = gr.Textbox( label="Text prompt for initial 3D generation", lines=1 ) gen_btn = gr.Button(value='Generate', scale=1) with gr.Column(): gr.HTML('Generated 3D objects') with gr.Row(): out_gen_1 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 1 (step 1)") out_gen_2 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 2 (step 1)") out_gen_3 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 3 (step 1)") out_gen_4 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 4 (step 1)") with gr.Column(scale=1): gr.HTML('Step 2: apply 3D editing with SHAP-EDITOR.') editing_choice = gr.Dropdown( ["Add a santa hat to it", "Make it look like made of gold", "Make the color of it look like rainbow", "Make it in cyberpunk style", "Make it wooden", "Make it look like make of lego"], value='Add a santa hat to it', multiselect=False, label="Editing effects", info="Select specific editing you want to apply!" ), apply_btn = gr.Button(value='Editing', scale=1) with gr.Column(scale=3): gr.HTML('Edited 3D objects') with gr.Row(): edited_1 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 1 (step 2)") edited_2 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 2 (step 2)") edited_3 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 3 (step 2)") edited_4 = gr.Model3D(clear_color=[0.0, 0.0, 0.0, 0.0], visible=True, label="3D Model 4 (step 2)") with gr.Accordion("Advanced Options", open=False): rand_seed = gr.Slider(minimum=0, maximum=1000, step=1, value=445, label="Random seed") gen_btn.click( fn=partial(generate_3d_with_shap_e, xm, diffusion, latent_model, device), inputs=[prompt, rand_seed, state], outputs=[out_gen_1, out_gen_2, out_gen_3, out_gen_4, state], queue=False) apply_btn.click( fn=partial(_3d_editing, xm, models, diffusion, initial_noise, noise_start_t, device), inputs=[ editing_choice[0], rand_seed, state ], outputs=[edited_1, edited_2, edited_3, edited_4, state], queue=True ) print("Generate examples...") with gr.Column(): gr.Examples( examples=[ [ "a corgi", "Make the color of it look like rainbow", 456, ], ["a penguin", "Make it look like make of lego", 214, ], ], inputs=[prompt, editing_choice[0], rand_seed], outputs=[out_gen_1, out_gen_2, out_gen_3, out_gen_4, edited_1, edited_2, edited_3, edited_4], fn=partial(optimize_all, xm, models, initial_noise, noise_start_t, diffusion, latent_model, device), cache_examples=True, ) demo.queue(max_size=10, api_open=False) demo.launch(share=True, show_api=False, show_error=True) if __name__ == '__main__': main()