Fix some issues with --gpu-only

jn-jairo · Jan 27, 2024 · 7f4725f · 7f4725f
1 parent fc196aa
commit 7f4725f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/comfy_extras/nodes_post_processing.py b/comfy_extras/nodes_post_processing.py
@@ -33,6 +33,7 @@ def INPUT_TYPES(s):
     CATEGORY = "image/postprocessing"
 
     def blend_images(self, image1: torch.Tensor, image2: torch.Tensor, blend_factor: float, blend_mode: str):
+        image2 = image2.to(image1.device)
         if image1.shape != image2.shape:
             image2 = image2.permute(0, 3, 1, 2)
             image2 = comfy.utils.common_upscale(image2, image1.shape[2], image1.shape[1], upscale_method='bicubic', crop='center')

diff --git a/comfy_extras/nodes_stable3d.py b/comfy_extras/nodes_stable3d.py
@@ -46,7 +46,7 @@ def encode(self, clip_vision, init_image, vae, width, height, batch_size, elevat
         encode_pixels = pixels[:,:,:,:3]
         t = vae.encode(encode_pixels)
         cam_embeds = camera_embeddings(elevation, azimuth)
-        cond = torch.cat([pooled, cam_embeds.repeat((pooled.shape[0], 1, 1))], dim=-1)
+        cond = torch.cat([pooled, cam_embeds.to(pooled.device).repeat((pooled.shape[0], 1, 1))], dim=-1)
 
         positive = [[cond, {"concat_latent_image": t}]]
         negative = [[torch.zeros_like(pooled), {"concat_latent_image": torch.zeros_like(t)}]]