Update app.py

app.py

@@ -11,11 +11,41 @@ os.system("git clone 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 utils.utils import  evaluate_in_chunks, scale_to_unit_sphere
 from network.model import StyleGAN2_3D
 
 
+def noise(batch_size, latent_dim, device):
+
+  return torch.randn(batch_size, latent_dim,device=device)
+
+
+def noise_list(batch_size, layers, latent_dim, device):
+  return [(noise(batch_size, latent_dim, device), layers)]
+
+
+def volume_noise(n, vol_size, device):
+  if device=="cuda":
+    return torch.FloatTensor(n, vol_size, vol_size, vol_size, 1).uniform_(0., 1.).cuda(device)
+  return torch.FloatTensor(n, vol_size, vol_size, vol_size, 1).uniform_(0., 1.)
+
+class StyleGAN2_3D_not_cuda(StyleGAN2_3D):
+
+    @torch.no_grad()
+    def generate_feature_volume(self, ema=False, trunc_psi=0.75):
+      latents = noise_list(
+          1, self.num_layers, self.latent_dim, device=self.device)
+      n = volume_noise(1, self.G_vol_size, device=self.device)
+      if ema:
+        generate_voxels = self.generate_truncated(
+            self.SE, self.GE, latents, n, trunc_psi)
+      else:
+        generate_voxels = self.generate_truncated(
+            self.S, self.G, latents, n, trunc_psi)
+
+      return generate_voxels
+
+
 cars=hf_hub_download("SerdarHelli/SDF-StyleGAN-3D", filename="cars.ckpt",revision="main")
 
 
@@ -25,7 +55,7 @@ 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 = StyleGAN2_3D_not_cuda.load_from_checkpoint(cars)
 model.eval()
 
 
@@ -46,9 +76,9 @@ def seed_all(seed):
 
 def change_model(ckpt_path):
     if device=="cuda":
-      model = StyleGAN2_3D.load_from_checkpoint(cars).cuda(0)
+      model = StyleGAN2_3D.load_from_checkpoint(ckpt_path).cuda(0)
     else:
-      model = StyleGAN2_3D.load_from_checkpoint(cars)
+      model = StyleGAN2_3D_not_cuda.load_from_checkpoint(ckpt_path)
     model.eval()