fix some errors. update code

.gitignore

@@ -23,3 +23,6 @@ scripts/research/
 .ipynb_checkpoints/
 _screenshots/
 flagged
+
+*.db
+gradio_queue.db

app.py

@@ -20,13 +20,20 @@ from huggingface_hub import hf_hub_download
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 port   = int(sys.argv[1]) if len(sys.argv) > 1 else 21111
 
+model_lists = {
+    'ffhq-512x512-basic':   dict(repo_id='facebook/stylenerf-ffhq-config-basic', filename='ffhq_512.pkl'),
+    'ffhq-256x256-basic':   dict(repo_id='facebook/stylenerf-ffhq-config-basic', filename='ffhq_256.pkl'), 
+    'ffhq-1024x1024-basic': dict(repo_id='facebook/stylenerf-ffhq-config-basic', filename='ffhq_1024.pkl'), 
+}
+model_names = [name for name in model_lists]
+
 
 def set_random_seed(seed):
     torch.manual_seed(seed)
     np.random.seed(seed)
 
 
-def get_camera_traj(model, pitch, yaw, fov=12, batch_size=1, model_name='FFHQ512'):
+def get_camera_traj(model, pitch, yaw, fov=12, batch_size=1, model_name=None):
     gen = model.synthesis
     range_u, range_v = gen.C.range_u, gen.C.range_v
     if not (('car' in model_name) or ('Car' in model_name)):  # TODO: hack, better option?
@@ -41,22 +48,10 @@ def get_camera_traj(model, pitch, yaw, fov=12, batch_size=1, model_name='FFHQ512
     return cam
 
 
-def check_name(model_name='FFHQ512'):
+def check_name(model_name):
     """Gets model by name."""
-    if model_name == 'FFHQ512':
-        network_pkl = hf_hub_download(repo_id='thomagram/stylenerf-ffhq-config-basic', filename='ffhq_512.pkl')
-
-    # TODO: checkpoint to be updated!
-    # elif model_name == 'FFHQ512v2':
-    #     network_pkl = "./pretrained/ffhq_512_eg3d.pkl"
-    # elif model_name == 'AFHQ512':
-    #     network_pkl = "./pretrained/afhq_512.pkl"
-    # elif model_name == 'MetFaces512':
-    #     network_pkl = "./pretrained/metfaces_512.pkl"
-    # elif model_name == 'CompCars256':
-    #     network_pkl = "./pretrained/cars_256.pkl"
-    # elif model_name == 'FFHQ1024':
-    #     network_pkl = "./pretrained/ffhq_1024.pkl"
+    if model_name in model_lists:
+        network_pkl = hf_hub_download(**model_lists[model_name])
     else:
         if os.path.isdir(model_name):
             network_pkl = sorted(glob.glob(model_name + '/*.pkl'))[-1]
@@ -85,7 +80,7 @@ def get_model(network_pkl, render_option=None):
     return G2, res, imgs
 
 
-global_states = list(get_model(check_name()))
+global_states = list(get_model(check_name(model_names[0])))
 wss  = [None, None]
 
 def proc_seed(history, seed):
@@ -98,7 +93,8 @@ def proc_seed(history, seed):
 def f_synthesis(model_name, model_find, render_option, early, trunc, seed1, seed2, mix1, mix2, yaw, pitch, roll, fov, history):
     history = history or {}
     seeds = []
-
+    trunc = trunc / 100
+    
     if model_find != "":
         model_name = model_find
 
@@ -124,7 +120,7 @@ def f_synthesis(model_name, model_find, render_option, early, trunc, seed1, seed
             set_random_seed(seed)
             z   = torch.from_numpy(np.random.RandomState(int(seed)).randn(1, model.z_dim).astype('float32')).to(device)
             ws  = model.mapping(z=z, c=None, truncation_psi=trunc)
-            img = model.get_final_output(styles=ws, camera_matrices=get_camera_traj(model, 0, 0), render_option=render_option)
+            img = model.get_final_output(styles=ws, camera_matrices=get_camera_traj(model, 0, 0, model_name=model_name), render_option=render_option)
             ws  = ws.detach().cpu().numpy()
             img = img[0].permute(1,2,0).detach().cpu().numpy()
 
@@ -178,26 +174,26 @@ def f_synthesis(model_name, model_find, render_option, early, trunc, seed1, seed
     image = (image * 255).astype('uint8')
     return image, history
 
-model_name = gr.inputs.Dropdown(['FFHQ512'])  #  'FFHQ512v2', 'AFHQ512', 'MetFaces512', 'CompCars256', 'FFHQ1024'
-model_find = gr.inputs.Textbox(label="checkpoint path", default="")
-render_option = gr.inputs.Textbox(label="rendering options", default='steps:40')
-trunc  = gr.inputs.Slider(default=0.7, maximum=1.0, minimum=0.0, label='truncation trick')
-seed1  = gr.inputs.Number(default=1, label="seed1")
-seed2  = gr.inputs.Number(default=9, label="seed2")
-mix1   = gr.inputs.Slider(minimum=0, maximum=1, default=0, label="linear mixing ratio (geometry)")
-mix2   = gr.inputs.Slider(minimum=0, maximum=1, default=0, label="linear mixing ratio (apparence)")
-early  = gr.inputs.Radio(['None', 'Normal Map', 'Gradient Map'], default='None', label='intermedia output')
-yaw    = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="yaw")
-pitch  = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="pitch")
-roll   = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="roll (optional, not suggested)")
-fov    = gr.inputs.Slider(minimum=9, maximum=15, default=12, label="fov")
+model_name = gr.inputs.Dropdown(model_names)
+model_find = gr.inputs.Textbox(label="Checkpoint path (folder or .pkl file)", default="")
+render_option = gr.inputs.Textbox(label="Additional rendering options", default='freeze_bg,steps:50')
+trunc  = gr.inputs.Slider(default=70, maximum=100, minimum=0, label='Truncation trick (%)')
+seed1  = gr.inputs.Number(default=1, label="Random seed1")
+seed2  = gr.inputs.Number(default=9, label="Random seed2")
+mix1   = gr.inputs.Slider(minimum=0, maximum=1, default=0, label="Linear mixing ratio (geometry)")
+mix2   = gr.inputs.Slider(minimum=0, maximum=1, default=0, label="Linear mixing ratio (apparence)")
+early  = gr.inputs.Radio(['None', 'Normal Map', 'Gradient Map'], default='None', label='Intermedia output')
+yaw    = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="Yaw")
+pitch  = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="Pitch")
+roll   = gr.inputs.Slider(minimum=-1, maximum=1, default=0, label="Roll (optional, not suggested for basic config)")
+fov    = gr.inputs.Slider(minimum=10, maximum=14, default=12, label="Fov")
 css = ".output-image, .input-image, .image-preview {height: 600px !important} "
 
 gr.Interface(fn=f_synthesis,
              inputs=[model_name, model_find, render_option, early, trunc, seed1, seed2, mix1, mix2, yaw, pitch, roll, fov, "state"],
-             title="Interctive Web Demo for StyleNeRF (ICLR 2022)",
-             description="Demo for ICLR 2022 Papaer: A Style-based 3D-Aware Generator for High-resolution Image Synthesis. Currently the demo runs on CPU only.",
+             title="Interactive Web Demo for StyleNeRF (ICLR 2022)",
+             description="StyleNeRF: A Style-based 3D-Aware Generator for High-resolution Image Synthesis. Currently the demo runs on CPU only.",
              outputs=["image", "state"],
              layout='unaligned',
-             css=css, theme='dark-huggingface',
+             css=css, theme='dark-seafoam',
              live=True).launch(enable_queue=True)