ShapE on Huggingface

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,13 @@
+---
+title: Shap E
+emoji: 🚀
+colorFrom: yellow
+colorTo: green
+sdk: gradio
+sdk_version: 3.28.2
+app_file: app.py
+pinned: true
+license: creativeml-openrail-m
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,268 @@
+import os
+import gradio as gr
+from PIL import Image
+import torch
+import matplotlib.pyplot as plt
+import imageio
+import numpy as np
+import math
+import argparse
+
+import torch
+import base64
+import io
+import os
+from typing import Union
+
+
+from shap_e.diffusion.sample import sample_latents
+from shap_e.diffusion.gaussian_diffusion import diffusion_from_config
+from shap_e.models.download import load_model, load_config
+from shap_e.util.notebooks import create_pan_cameras, decode_latent_images
+
+from shap_e.models.nn.camera import DifferentiableCameraBatch, DifferentiableProjectiveCamera
+from shap_e.models.transmitter.base import Transmitter, VectorDecoder
+from shap_e.util.collections import AttrDict
+
+import trimesh
+
+
+state = ""
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+css = '''
+    .instruction{position: absolute; top: 0;right: 0;margin-top: 0px !important}
+    .arrow{position: absolute;top: 0;right: -110px;margin-top: -8px !important}
+    #component-4, #component-3, #component-10{min-height: 0}
+    .duplicate-button img{margin: 0}
+'''
+
+def set_state(s):
+    print(s)
+    global state
+    state = s
+
+def get_state():
+    return state
+
+@torch.no_grad()
+def decode_latent_images_foo(
+    xm: Union[Transmitter, VectorDecoder],
+    latent: torch.Tensor,
+    cameras: DifferentiableCameraBatch,
+    rendering_mode: str = "stf",
+):
+    decoded = xm.renderer.render_views(
+        AttrDict(cameras=cameras),
+        params=(xm.encoder if isinstance(xm, Transmitter) else xm).bottleneck_to_params(
+            latent[None]
+        ),
+        options=AttrDict(rendering_mode=rendering_mode, render_with_direction=False),
+    )
+    return decoded
+
+def generate_3D(input, grid_size=64):
+    set_state('Entered generate function...')
+
+    # if input is a string, it's a text prompt
+    xm = load_model('transmitter', device=device)
+    diffusion = diffusion_from_config(load_config('diffusion'))
+    batch_size = 4
+
+    print(type(input))
+
+    if isinstance(input, Image.Image):
+        input = prepare_img(input)
+        model = load_model('image300M', device=device)
+        guidance_scale = 3.0
+        model_kwargs = dict(images=[input] * batch_size)
+    else:
+        model = load_model('text300M', device=device)
+        guidance_scale = 15.0
+        model_kwargs = dict(texts=[input] * batch_size)
+
+    print(input)
+
+    latents = sample_latents(
+        batch_size=batch_size,
+        model=model,
+        diffusion=diffusion,
+        guidance_scale=guidance_scale,
+        model_kwargs=model_kwargs,
+        progress=True,
+        clip_denoised=True,
+        use_fp16=True,
+        use_karras=True,
+        karras_steps=64,
+        sigma_min=1e-3,
+        sigma_max=160,
+        s_churn=0,
+    )
+
+    render_mode = 'stf' # you can change this to 'stf'
+    size = grid_size # this is the size of the renders; higher values take longer to render.
+
+    cameras = create_pan_cameras(size, device)
+
+    x=decode_latent_images_foo(xm, latents[0], cameras, rendering_mode=render_mode)
+    mesh=x['meshes'][0]
+    rm=x['raw_meshes'][0]
+
+    rm.vertex_channels["R"]=mesh.vertex_colors[:,0]
+    rm.vertex_channels["G"]=mesh.vertex_colors[:,1]
+    rm.vertex_channels["B"]=mesh.vertex_colors[:,2]
+
+    tm=rm.tri_mesh()
+
+    with open("output/mesh.ply",'wb') as f:
+        tm.write_ply(f)
+
+
+    set_state('Converting to point cloud...')
+    # pc = sampler.output_to_point_clouds(samples)[0]
+
+    set_state('Converting to mesh...')
+    # save_ply(pc, 'output/mesh.ply', grid_size)
+
+    set_state('')
+
+    images = decode_latent_images(xm, latents[0], cameras, rendering_mode=render_mode)
+
+    # Convert images to gif
+    os.makedirs('output', exist_ok=True)
+    images[0].save(f'/tmp/mesh.mp4', save_all=True, append_images=images[1:], duration=100, loop=0)
+
+
+    return ply_to_glb('/tmp/mesh.ply', '/tmp/mesh.glb'), '/tmp/mesh.mp4', gr.update(value=['/tmp/mesh.glb', '/tmp/mesh.ply'], visible=True)
+
+def prepare_img(img):
+
+    w, h = img.size
+    if w > h:
+        img = img.crop((w - h) / 2, 0, w - (w - h) / 2, h)
+    else:
+        img = img.crop((0, (h - w) / 2, w, h - (h - w) / 2))
+
+    # resize to 256x256
+    img = img.resize((256, 256))
+
+    return img
+
+
+def ply_to_glb(ply_file, glb_file):
+    mesh = trimesh.load(ply_file)
+
+    # Save the mesh as a glb file using Trimesh
+    mesh.export(glb_file, file_type='glb')
+
+    return glb_file
+
+
+# def save_ply(pc, file_name, grid_size):
+#     set_state('Creating SDF model...')
+#     sdf_name = 'sdf'
+#     sdf_model = model_from_config(MODEL_CONFIGS[sdf_name], device)
+#     sdf_model.eval()
+
+#     set_state('Loading SDF model...')
+#     sdf_model.load_state_dict(load_checkpoint(sdf_name, device))
+
+#     # Produce a mesh (with vertex colors)
+#     mesh = marching_cubes_mesh(
+#         pc=pc,
+#         model=sdf_model,
+#         batch_size=4096,
+#         grid_size=grid_size, # increase to 128 for resolution used in evals
+#         progress=True,
+#     )
+
+#     # Write the mesh to a PLY file to import into some other program.
+#     with open(file_name, 'wb') as f:
+#         mesh.write_ply(f)
+
+
+block = gr.Blocks().queue(max_size=250, concurrency_count=6)
+with block:
+    with gr.Box():
+        if(not torch.cuda.is_available()):
+            top_description = gr.HTML(f'''
+                <div style="text-align: center; max-width: 650px; margin: 0 auto;">
+                <div>
+                    <img class="logo" src="file/images/mirage.png" alt="Mirage Logo"
+                        style="margin: auto; max-width: 7rem;">
+                    <br />
+                    <h1 style="font-weight: 900; font-size: 2.5rem;">
+                    Shap-E Web UI
+                    </h1>
+                    <br />
+                    <a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/MirageML/point-e?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a>
+                </div>
+                <br />
+                <p style="margin-bottom: 10px; font-size: 94%">
+                Generate 3D Assets in 1 minute with a prompt or image!
+                Based on the <a href="https://github.com/openai/shap-e">Shap-E</a> implementation
+                </p>
+                <br />
+                <p>There's only one step left before you can train your model: <a href="https://huggingface.co/spaces/{os.environ['SPACE_ID']}/settings" style="text-decoration: underline" target="_blank">attribute a <b>T4 GPU</b> to it (via the Settings tab)</a> and run the training below. Other GPUs are not compatible for now. You will be billed by the minute from when you activate the GPU until when it is turned it off.</p>
+                </div>
+            ''')
+        else:
+            top_description = gr.HTML(f'''
+                    <div style="text-align: center; max-width: 650px; margin: 0 auto;">
+                    <div>
+                        <img class="logo" src="file/images/mirage.png" alt="Mirage Logo"
+                            style="margin: auto; max-width: 7rem;">
+                        <br />
+                        <h1 style="font-weight: 900; font-size: 2.5rem;">
+                        Point-E Web UI
+                        </h1>
+                        <br />
+                        <a class="duplicate-button" style="display:inline-block" target="_blank" href="https://huggingface.co/spaces/MirageML/point-e?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a>
+                    </div>
+                    <br />
+                    <p style="margin-bottom: 10px; font-size: 94%">
+                    Generate 3D Assets in 1 minute with a prompt or image!
+                    Based on the <a href="https://github.com/openai/shap-e">Shap-E</a> implementation
+                    </p>
+                    </div>
+                ''')
+    with gr.Row():
+        with gr.Column():
+            with gr.Tab("Text to 3D"):
+                gr.Markdown("Uses Stable Diffusion to create an image from the prompt.")
+                prompt = gr.Textbox(label="Prompt", placeholder="A HD photo of a Corgi")
+                text_button = gr.Button(label="Generate")
+
+            with gr.Tab("Image to 3D"):
+                gr.Markdown("Best results with images of objects on an empty background.")
+                input_image = gr.Image(label="Image")
+                img_button = gr.Button(label="Generate")
+
+            # with gr.Accordion("Advanced options", open=False):
+            #     model = gr.Radio(["base40M", "base300M", "base1B"], label="Model", value="base1B")
+            #     scale = gr.Slider(
+            #         label="Guidance Scale", minimum=1.0, maximum=10.0, value=3.0, step=0.1
+            #     )
+
+        with gr.Column():
+            model_gif = gr.Model3D(label="3D Model GIF")
+            # btn_pc_to_obj = gr.Button(value="Convert to OBJ", visible=False)
+            model_3d = gr.Model3D(value=None)
+            file_out = gr.File(label="Files", visible=False)
+
+    if torch.cuda.is_available():
+        gr.Examples(
+            examples=[
+                ["images/pumpkin.png"],
+                ["images/fantasy_world.png"],
+            ],
+            inputs=[input_image],
+            outputs=[model_3d, model_gif, file_out],
+            fn=generate_3D,
+            cache_examples=True
+        )
+
+    img_button.click(fn=generate_3D, inputs=[input_image], outputs=[model_3d, model_gif, file_out])
+    text_button.click(fn=generate_3D, inputs=[prompt], outputs=[model_3d, model_gif, file_out])
+
+block.launch(show_api=False)

requirements.txt

@@ -0,0 +1,9 @@
+git+https://github.com/openai/shap-e@main
+trimesh
+diffusers
+transformers
+accelerate
+imageio[ffmpeg]
+imageio[pyav]
+blobfile
+ipywidgets

shap_e_model_cache/ViT-L-14.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b8cca3fd41ae0c99ba7e8951adf17d267cdb84cd88be6f7c2e0eca1737a03836
+size 932768134

shap_e_model_cache/diffusion_config.yaml

@@ -0,0 +1,3 @@
+mean_type: x_start
+schedule: exp
+timesteps: 1024

shap_e_model_cache/image_cond.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb8072c64bbbcf6910488814d212227de5db291780d4ea99c6152f9346cf12aa
+size 1263925407

shap_e_model_cache/image_cond_config.yaml

@@ -0,0 +1,13 @@
+d_latent: 1048576
+inner:
+  cond_drop_prob: 0.1
+  heads: 8
+  init_scale: 0.25
+  layers: 24
+  name: CLIPImageGridPointDiffusionTransformer
+  pos_emb_init_scale: 0.05
+  time_token_cond: true
+  use_pos_emb: true
+  width: 1024
+latent_ctx: 1024
+name: SplitVectorDiffusion

shap_e_model_cache/text_cond.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e6b4fa599a7b3c3b16c222d5f5fe56f9db9289ff0b6575fbe5c11bc97106aad4
+size 1262868003

shap_e_model_cache/text_cond_config.yaml

@@ -0,0 +1,14 @@
+d_latent: 1048576
+inner:
+  cond_drop_prob: 0.1
+  heads: 16
+  init_scale: 0.25
+  layers: 24
+  name: CLIPImagePointDiffusionTransformer
+  pos_emb_init_scale: 0.05
+  time_token_cond: true
+  token_cond: true
+  use_pos_emb: true
+  width: 1024
+latent_ctx: 1024
+name: SplitVectorDiffusion

shap_e_model_cache/transmitter.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:af02a0b85a8abdfb3919584b63c540ba175f6ad4790f574a7fef4617e5acdc3b
+size 1776048343

shap_e_model_cache/transmitter_config.yaml

@@ -0,0 +1,80 @@
+encoder:
+  cross_attention_dataset: pcl_and_multiview_pcl
+  d_latent: 1024
+  data_ctx: 1024
+  fps_method: first
+  heads: 8
+  init_scale: 0.25
+  inner_batch_size:
+  - 16384
+  - 20
+  input_channels: 6
+  latent_bottleneck:
+    diffusion:
+      schedule: inv_parabola
+      schedule_args:
+        power: 5.0
+      timesteps: 1024
+    diffusion_prob: 0.1
+    name: clamp_diffusion_noise
+  layers: 12
+  max_depth: 9.0
+  max_unrolls: 1
+  min_unrolls: 1
+  name: PointCloudPerceiverChannelsEncoder
+  params_proj:
+    init_scale: 1.0
+    learned_scale: 0.0625
+    name: channels
+    use_ln: true
+  patch_size: 8
+  pointconv_hidden:
+  - 1024
+  - 1024
+  pointconv_padding_mode: circular
+  pointconv_patch_size: 8
+  pointconv_samples: 64
+  pointconv_stride: 4
+  pos_emb: nerf
+  use_depth: true
+  use_pointconv: true
+  width: 1024
+name: Transmitter
+renderer:
+  grid_size: 128
+  n_coarse_samples: 64
+  n_fine_samples: 128
+  name: NeRSTFRenderer
+  nerstf:
+    activation: swish
+    d_hidden: 256
+    density_activation: relu
+    init_scale: 0.25
+    initial_density_bias: 0.1
+    insert_direction_at: 4
+    meta_bias: false
+    meta_parameters: true
+    n_hidden_layers: 6
+    n_meta_layers: 4
+    name: MLPNeRSTFModel
+    posenc_version: nerf
+    separate_coarse_channels: true
+    separate_nerf_channels: true
+    trainable_meta: false
+  separate_shared_samples: true
+  void:
+    background:
+    - 0
+    - 0
+    - 0
+    name: VoidNeRFModel
+  volume:
+    bbox_max:
+    - 1.0
+    - 1.0
+    - 1.0
+    bbox_min:
+    - -1.0
+    - -1.0
+    - -1.0
+    name: BoundingBoxVolume