Hugging Face's logo

Spaces:
SerdarHelli
/
SDF-StyleGan-3D
Running

App Files Community
SDF-StyleGan-3D
File size: 3,793 Bytes 
a7808a2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
a589fb1
a7808a2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
47dddc9
a7808a2
47dddc9
a7808a2
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
47dddc9
a7808a2
 
 
 
47dddc9
 
a7808a2
 
 
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
 
import os
import gradio as gr
import plotly.graph_objects as go
import sys
import torch
from huggingface_hub import hf_hub_download
import numpy as np
import random

os.system("https://github.com/Zhengxinyang/SDF-StyleGAN.git")
sys.path.append("SDF-StyleGAN")

#Codes reference : https://github.com/Zhengxinyang/SDF-StyleGAN

from utils.utils import noise, evaluate_in_chunks, scale_to_unit_sphere, volume_noise, process_sdf, linear_slerp
from network.model import StyleGAN2_3D


cars=hf_hub_download("SerdarHelli/SDF-StyleGAN-3D", filename="cars.ckpt",revision="main")



#default model
device='cuda' if torch.cuda.is_available() else 'cpu'
if device=="cuda":
  model = StyleGAN2_3D.load_from_checkpoint(cars).cuda(0)
else:
  model = StyleGAN2_3D.load_from_checkpoint(cars)
model.eval()


models={"Car":cars,
        "Airplane":"./planes.ckpt",
        "Chair":"./chairs.ckpt",
        "Rifle":"./rifles.ckpt",
        "Table":"./tables.ckpt"
}


def seed_all(seed):

    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)


def change_model(ckpt_path):
    if device=="cuda":
      model = StyleGAN2_3D.load_from_checkpoint(cars).cuda(0)
    else:
      model = StyleGAN2_3D.load_from_checkpoint(cars)
    model.eval()


def predict(seed,trunc_psi):
  if seed==None:
    seed=777 
  seed_all(seed)
  if trunc_psi==None:
    trunc_psi=1

  z = noise(100000, model.latent_dim, device=model.device)
  samples = evaluate_in_chunks(1000, model.SE, z)
  model.av = torch.mean(samples, dim=0, keepdim=True)

  mesh = model.generate_mesh(
      ema=True, mc_vol_size=64, level=-0.015, trunc_psi=trunc_psi)
  mesh = scale_to_unit_sphere(mesh)
  mesh.export("/content/asdads.obj")
  x=np.asarray(mesh.vertices).T[0]
  y=np.asarray(mesh.vertices).T[1]
  z=np.asarray(mesh.vertices).T[2]

  i=np.asarray(mesh.faces).T[0]
  j=np.asarray(mesh.faces).T[1]
  k=np.asarray(mesh.faces).T[2]

  return x,y,z,i,j,k

def generate(seed,model_name,trunc_psi):

    change_model(models[model_name])
    x,y,z,i,j,k=predict(seed,trunc_psi)


    fig = go.Figure(go.Mesh3d(x=x, y=y, z=z, 
                    i=i, j=j, k=k, 
                    colorscale="Viridis",
                  colorbar_len=0.75,
                  flatshading=True,
                  lighting=dict(ambient=0.5,
                                diffuse=1,
                                fresnel=4,        
                                specular=0.5,
                                roughness=0.05,
                                facenormalsepsilon=0,
                                vertexnormalsepsilon=0),
                  lightposition=dict(x=100,
                                    y=100,
                                    z=1000)))
    return fig
    
markdown=f'''
  # SDF-StyleGAN: Implicit SDF-Based StyleGAN for 3D Shape Generation

  
  [The space demo for the SGP 2022 paper "SDF-StyleGAN: Implicit SDF-Based StyleGAN for 3D Shape Generation".](https://arxiv.org/abs/2206.12055)
  
  [For the official implementation.](https://github.com/Zhengxinyang/SDF-StyleGAN)
  ### Future Work based on interest
  - Adding new models for new type objects
  - New Customization 
  
  
  It is running on {device}
  
'''
with gr.Blocks() as demo:
    with gr.Column():
        with gr.Row():
            gr.Markdown(markdown)
        with gr.Row():
            seed = gr.Slider( minimum=0, maximum=2**16,label='Seed')
            model_name=gr.Dropdown(choices=["Car","Airplane","Chair","Rifle","Table"],label="Choose Model Type")
            trunc_psi = gr.Slider( minimum=0, maximum=2,label='Truncate PSI')

        btn = gr.Button(value="Generate")
        mesh = gr.Plot()
    demo.load(generate, [seed,model_name,trunc_psi], mesh)
    btn.click(generate, [seed,model_name,trunc_psi], mesh)

demo.launch(debug=True)