fix

app.py

@@ -37,7 +37,7 @@ device= 'cpu'
 ffhq_model_path = hf_hub_download(repo_id="bankholdup/stylegan_petbreeder", filename="e4e_ffhq512.pt")
 
 ffhq_ckpt = torch.load(ffhq_model_path, map_location='cpu')
-ffhq_latent_avg = ffhq_ckpt['latent_avg'].to('cuda:0')
+ffhq_latent_avg = ffhq_ckpt['latent_avg'].to(device)
 ffhq_opts = ffhq_ckpt['opts']
 ffhq_opts['checkpoint_path'] = ffhq_model_path
 ffhq_opts= Namespace(**ffhq_opts)
@@ -58,7 +58,7 @@ clear_output()
 dog_model_path = hf_hub_download(repo_id="bankholdup/stylegan_petbreeder", filename="e4e_ffhq512_dog.pt")
 
 dog_ckpt = torch.load(dog_model_path, map_location='cpu')
-dog_latent_avg = dog_ckpt['latent_avg'].to('cuda:0')
+dog_latent_avg = dog_ckpt['latent_avg'].to(device)
 dog_opts = dog_ckpt['opts']
 dog_opts['checkpoint_path'] = dog_model_path
 dog_opts= Namespace(**dog_opts)
@@ -79,7 +79,7 @@ clear_output()
 cat_model_path = hf_hub_download(repo_id="bankholdup/stylegan_petbreeder", filename="e4e_ffhq512_cat.pt")
 
 cat_ckpt = torch.load(cat_model_path, map_location='cpu')
-cat_latent_avg = cat_ckpt['latent_avg'].to('cuda:0')
+cat_latent_avg = cat_ckpt['latent_avg'].to(device)
 cat_opts = cat_ckpt['opts']
 cat_opts['checkpoint_path'] = cat_model_path
 cat_opts= Namespace(**cat_opts)
@@ -114,13 +114,13 @@ def inference(img):
     img.save('out.jpg')  
     aligned_face = align_face('out.jpg')
 	
-    ffhq_codes = ffhq_encoder(aligned_face.unsqueeze(0).to("cuda").float())
+    ffhq_codes = ffhq_encoder(aligned_face.unsqueeze(0).to(device).float())
     ffhq_codes = ffhq_codes + ffhq_latent_avg.repeat(ffhq_codes.shape[0], 1, 1)
 
-    cat_codes = cat_encoder(aligned_face.unsqueeze(0).to("cuda").float())
+    cat_codes = cat_encoder(aligned_face.unsqueeze(0).to(device).float())
     cat_codes = cat_codes + ffhq_latent_avg.repeat(cat_codes.shape[0], 1, 1)
 
-    dog_codes = dog_encoder(aligned_face.unsqueeze(0).to("cuda").float())
+    dog_codes = dog_encoder(aligned_face.unsqueeze(0).to(device).float())
     dog_codes = dog_codes + ffhq_latent_avg.repeat(dog_codes.shape[0], 1, 1)
         
     animal = "cat"