Update app.py

app.py

@@ -52,11 +52,7 @@ cars=hf_hub_download("SerdarHelli/SDF-StyleGAN-3D", filename="cars.ckpt",revisio
 
 #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_not_cuda.load_from_checkpoint(cars)
-model.eval()
+
 
 
 models={"Car":cars,
@@ -74,15 +70,9 @@ def seed_all(seed):
     random.seed(seed)
 
 
-def change_model(ckpt_path):
-    if device=="cuda":
-      model = StyleGAN2_3D.load_from_checkpoint(ckpt_path).cuda(0)
-    else:
-      model = StyleGAN2_3D_not_cuda.load_from_checkpoint(ckpt_path)
-    model.eval()
 
 
-def predict(seed,trunc_psi):
+def predict(seed,model,trunc_psi):
   if seed==None:
     seed=777 
   seed_all(seed)
@@ -107,9 +97,22 @@ def predict(seed,trunc_psi):
   return x,y,z,i,j,k
 
 def generate(seed,model_name,trunc_psi):
-    if model_name:
-        change_model(models[model_name])
-    x,y,z,i,j,k=predict(seed,trunc_psi)
+    print(model_name)
+    try :
+      ckpt=models[model_name]
+    except KeyError:
+      ckpt=cars
+
+
+    if device=="cuda":
+        model = StyleGAN2_3D.load_from_checkpoint(ckpt).cuda(0)
+    else:
+        model = StyleGAN2_3D_not_cuda.load_from_checkpoint(ckpt)
+    model.eval()
+
+
+
+    x,y,z,i,j,k=predict(seed,model,trunc_psi)
 
 
     fig = go.Figure(go.Mesh3d(x=x, y=y, z=z, 
@@ -131,7 +134,6 @@ def generate(seed,model_name,trunc_psi):
     
 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)
   
@@ -150,7 +152,7 @@ with gr.Blocks() as demo:
             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")
+            model_name=gr.Dropdown(choices=["Car","Airplane"],label="Choose Model Type")
             trunc_psi = gr.Slider( minimum=0, maximum=2,label='Truncate PSI')
 
         btn = gr.Button(value="Generate")