app.py

import os
import gradio as gr
from PIL import Image
import torch
import matplotlib.pyplot as plt
import imageio
import numpy as np
import argparse

from point_e.diffusion.configs import DIFFUSION_CONFIGS, diffusion_from_config
from point_e.diffusion.sampler import PointCloudSampler
from point_e.models.download import load_checkpoint
from point_e.models.configs import MODEL_CONFIGS, model_from_config
from point_e.util.plotting import plot_point_cloud
from point_e.util.pc_to_mesh import marching_cubes_mesh

from diffusers import StableDiffusionPipeline

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

def load_img2mesh_model(model_name):
    set_state(f'Creating img2mesh model {model_name}...')
    i2m_name = model_name
    i2m_model = model_from_config(MODEL_CONFIGS[i2m_name], device)
    i2m_model.eval()
    base_diffusion_i2m = diffusion_from_config(DIFFUSION_CONFIGS[i2m_name])

    set_state(f'Downloading img2mesh checkpoint {model_name}...')
    i2m_model.load_state_dict(load_checkpoint(i2m_name, device))

    return i2m_model, base_diffusion_i2m



def get_sampler(model_name, txt2obj, guidance_scale):
    if txt2obj:
        set_state('Creating txt2mesh model...')
        t2m_name = 'base40M-textvec'
        t2m_model = model_from_config(MODEL_CONFIGS[t2m_name], device)
        t2m_model.eval()
        base_diffusion_t2m = diffusion_from_config(DIFFUSION_CONFIGS[t2m_name])

        set_state('Downloading txt2mesh checkpoint...')
        t2m_model.load_state_dict(load_checkpoint(t2m_name, device))
    else:
        i2m_model, base_diffusion_i2m = load_img2mesh_model(model_name)

    set_state('Creating upsample model...')
    upsampler_model = model_from_config(MODEL_CONFIGS['upsample'], device)
    upsampler_model.eval()
    upsampler_diffusion = diffusion_from_config(DIFFUSION_CONFIGS['upsample'])

    set_state('Downloading upsampler checkpoint...')
    upsampler_model.load_state_dict(load_checkpoint('upsample', device))

    return PointCloudSampler(
            device=device,
            models=[t2m_model if txt2obj else i2m_model, upsampler_model],
            diffusions=[base_diffusion_t2m if txt2obj else base_diffusion_i2m, upsampler_diffusion],
            num_points=[1024, 4096 - 1024],
            aux_channels=['R', 'G', 'B'],
            guidance_scale=[guidance_scale, 0.0 if txt2obj else guidance_scale],
            model_kwargs_key_filter=('texts', '') if txt2obj else ("*",)
        )

def generate_txt2img(prompt):
    pipe = StableDiffusionPipeline.from_pretrained("point_e_model_cache/stable-diffusion-2-1", torch_dtype=torch.float16)
    pipe = pipe.to("cuda")
    image = pipe(prompt).images[0]

    return image

def generate_3D(input, model_name='base1B', guidance_scale=3.0, grid_size=128):
    set_state('Entered generate function...')

    # try:
    #     input = Image.fromarray(input)
    # except:
    #     img = generate_txt2img(input)
    #     img.save('/tmp/img.png')
    #     input = Image.open('/tmp/img.png')

    if isinstance(input, Image.Image):
        input = prepare_img(input)

    # if input is a string, it's a text prompt
    sampler = get_sampler(model_name, txt2obj=True if isinstance(input, str) else False, guidance_scale=guidance_scale)

    # Produce a sample from the model.
    set_state('Sampling...')
    samples = None
    kw_args = dict(texts=[input]) if isinstance(input, str) else dict(images=[input])
    for x in sampler.sample_batch_progressive(batch_size=1, model_kwargs=kw_args):
        samples = x

    set_state('Converting to point cloud...')
    pc = sampler.output_to_point_clouds(samples)[0]

    set_state('Converting to mesh...')
    save_ply(pc, '/tmp/mesh.ply', grid_size)

    set_state('')

    return ply_to_glb('/tmp/mesh.ply', '/tmp/mesh.glb'), create_gif(pc), 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)

def create_gif(pc):
    fig = plt.figure(facecolor='black', figsize=(4, 4))
    ax = fig.add_subplot(111, projection='3d', facecolor='black')
    fixed_bounds=((-0.75, -0.75, -0.75),(0.75, 0.75, 0.75))

    # Create an empty list to store the frames
    frames = []

    # Create a loop to generate the frames for the GIF
    for angle in range(0, 360, 4):
        # Clear the plot and plot the point cloud
        ax.clear()
        color_args = np.stack(
            [pc.channels["R"], pc.channels["G"], pc.channels["B"]], axis=-1
        )
        c = pc.coords


        ax.scatter(c[:, 0], c[:, 1], c[:, 2], c=color_args)

        # Set the viewpoint for the plot
        ax.view_init(elev=10, azim=angle)

        # Turn off the axis labels and ticks
        ax.axis('off')
        ax.set_xlim3d(fixed_bounds[0][0], fixed_bounds[1][0])
        ax.set_ylim3d(fixed_bounds[0][1], fixed_bounds[1][1])
        ax.set_zlim3d(fixed_bounds[0][2], fixed_bounds[1][2])

        # Draw the figure to update the image data
        fig.canvas.draw()

        # Save the plot as a frame for the GIF
        frame = np.array(fig.canvas.renderer.buffer_rgba())
        w, h = frame.shape[0], frame.shape[1]
        i = int(round((h - int(h*0.6)) / 2.))
        frame = frame[i:i + int(h*0.6),i:i + int(h*0.6)]
        frames.append(frame)

    # Save the GIF using imageio
    imageio.mimsave('/tmp/pointcloud.mp4', frames, fps=30)
    return '/tmp/pointcloud.mp4'

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;">
                    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 2 minutes with a prompt or image!
                Based on the <a href="https://github.com/openai/point-e">Point-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 2 minutes with a prompt or image!
                    Based on the <a href="https://github.com/openai/point-e">Point-E</a> implementation
                    </p>
                    </div>
                ''')
    with gr.Row():
        with gr.Column():
            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.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.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.Video(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, model, scale], outputs=[model_3d, model_gif, file_out])
    text_button.click(fn=generate_3D, inputs=[prompt, model, scale], outputs=[model_3d, model_gif, file_out])

block.launch(show_api=False)