update cluster plot cuda

app.py

@@ -497,6 +497,8 @@ def ncut_run(
         if not video_output:
             start = time.time()
             h, w = features.shape[1], features.shape[2]
+            if torch.cuda.is_available():
+                images = images.cuda()
             _images = reverse_transform_image(images, stablediffusion="stable" in model_name.lower())
             cluster_images = make_cluster_plot(eigvecs, _images, h=h, w=w)
             logging_str += f"Plot time: {time.time() - start:.2f}s\n"
@@ -602,9 +604,9 @@ def reverse_transform_image(image, stablediffusion=False):
     if stablediffusion:
         image = (image + 1) / 2
     else:
-        mean = [0.485, 0.456, 0.406]
-        std = [0.229, 0.224, 0.225]
-        image = image * torch.tensor(std).view(3, 1, 1) + torch.tensor(mean).view(3, 1, 1)
+        mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1).to(image.device)
+        std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1).to(image.device)
+        image = image * std + mean
     image = torch.clamp(image, 0, 1)
     return image
 
@@ -1168,7 +1170,7 @@ with demo:
                 
             with gr.Column(scale=5, min_width=200):
                 output_gallery = make_output_images_section()
-                cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=False, elem_id="clusters", columns=[2], rows=[1], object_fit="contain", height="auto", show_share_button=True, preview=True)
+                cluster_gallery = gr.Gallery(value=[], label="Clusters", show_label=False, elem_id="clusters", columns=[2], rows=[1], object_fit="contain", height=450, show_share_button=True, preview=True)
                 [
                     model_dropdown, layer_slider, node_type_dropdown, num_eig_slider, 
                     affinity_focal_gamma_slider, num_sample_ncut_slider, knn_ncut_slider,