Fix Gradio slider update and CUDA device issues for all attributes

app.py

@@ -206,7 +206,7 @@ def build_ui() -> gr.Blocks:
             lo, hi = recommended_range(name)
             # Keep current value within new bounds
             new_val = max(lo, min(hi, strength.value if hasattr(strength, "value") else 0))
-            return gr.Slider.update(minimum=lo, maximum=hi, value=new_val)
+            return gr.Slider(minimum=lo, maximum=hi, value=new_val)
 
         attr.change(_on_attr_change, inputs=attr, outputs=strength)
 

editings/ganspace.py

@@ -6,16 +6,17 @@ def edit(latents, pca, edit_directions):
     for latent in latents:
         for pca_idx, start, end, strength in edit_directions:
             delta = get_delta(pca, latent, pca_idx, strength)
-            delta_padded = torch.zeros(latent.shape).to("cuda")
+            delta_padded = torch.zeros(latent.shape).to(latent.device)
             delta_padded[start:end] += delta.repeat(end - start, 1)
             edit_latents.append(latent + delta_padded)
     return torch.stack(edit_latents)
 
 
 def get_delta(pca, latent, idx, strength):
-    w_centered = latent - pca["mean"].to("cuda")
-    lat_comp = pca["comp"].to("cuda")
-    lat_std = pca["std"].to("cuda")
+    device = latent.device
+    w_centered = latent - pca["mean"].to(device)
+    lat_comp = pca["comp"].to(device)
+    lat_std = pca["std"].to(device)
     w_coord = (
         torch.sum(w_centered[0].reshape(-1) * lat_comp[idx].reshape(-1)) / lat_std[idx]
     )
@@ -26,7 +27,7 @@ def get_delta(pca, latent, idx, strength):
 def edit_latent(latent, pca, edit_direction):
     pca_idx, start, end, strength = edit_direction
     delta = get_delta(pca, latent, pca_idx, strength)
-    delta_padded = torch.zeros(latent.shape).to("cuda")
+    delta_padded = torch.zeros(latent.shape).to(latent.device)
     delta_padded[start:end] += delta.repeat(end - start, 1)
     edit_latent = latent + delta_padded
     return edit_latent