File size: 9,282 Bytes
e43a382
 
 
 
 
 
 
 
 
a698138
e43a382
 
 
 
 
 
 
 
 
 
 
fad0371
e43a382
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a698138
 
 
 
 
 
 
 
e43a382
 
 
 
 
 
 
 
f845bf6
 
e43a382
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fad0371
 
e43a382
 
 
 
 
 
 
 
 
 
 
 
 
a698138
e43a382
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0e30a9a
f845bf6
0e30a9a
f845bf6
e43a382
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
0e30a9a
e43a382
 
0e30a9a
f845bf6
0e30a9a
f845bf6
e43a382
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a698138
 
e43a382
a698138
e43a382
 
 
 
f845bf6
 
a698138
 
f845bf6
a698138
f845bf6
 
 
 
e43a382
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
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 tempfile

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, decode_latent_mesh

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}
'''

def set_state(s):
    print(s)
    global state
    state = s

def get_state():
    return state

def to_video(frames: list[Image.Image], fps: int = 5) -> str:
        out_file = tempfile.NamedTemporaryFile(suffix='.mp4', delete=False)
        writer = imageio.get_writer(out_file.name, format='FFMPEG', fps=fps)
        for frame in frames:
            writer.append_data(np.asarray(frame))
        writer.close()
        return out_file.name

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

    if isinstance(input, np.ndarray):
        input = Image.fromarray(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)

    with open(f'/tmp/mesh.ply', 'wb') as f:
        decode_latent_mesh(xm, latents[0]).tri_mesh().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)


    return ply_to_glb('/tmp/mesh.ply', '/tmp/mesh.glb'), to_video(images), 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>
                </div>
                <h3 style="font-weight: 900; font-size: 1.5rem;">
                If the Queue is Too Long, <a href="https://app.mirageml.com/generate" style="text-decoration: underline" target="_blank">Try it on Mirage</a>!
                </h3>
                <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;">
                        Shap-E Web UI
                        </h1>
                    </div>
                    <h3 style="font-weight: 900; font-size: 1.5rem;">
                    If the Queue is Too Long, <a href="https://app.mirageml.com/generate" style="text-decoration: underline" target="_blank">Try it on Mirage</a>!
                    </h3>
                    <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=[
                ["a shark"],
                ["an avocado"],
            ],
            inputs=[prompt],
            outputs=[model_3d, model_gif, file_out],
            fn=generate_3D,
            cache_examples=True
        )
        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)