298977887/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-01-11 02:15:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

Only use events for devices supporting nonblocking

Browse Source
This commit is contained in:
Austin Mroz 2024-12-23 16:43:43 -06:00
parent 813b8df872
commit ce5afecc36
No known key found for this signature in database
1 changed files with 13 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
latent_preview.py
View File

@ -18,16 +18,21 @@ class LatentPreviewer:
# A previous preview is still being processed
return None
preview_tensor = self.decode_latent_to_preview(x0)
if comfy.model_management.device_supports_non_blocking(preview_tensor.device):
latents_ubyte = (((preview_tensor + 1.0) / 2.0).clamp(0, 1) # change scale from -1..1 to 0..1
.mul(0xFF) # to 0..255
).to(device="cpu", dtype=torch.uint8, non_blocking=True)
latents_rgbx = torch.zeros(latents_ubyte.shape[:2] + (4,), device="cpu", dtype=torch.uint8)
latents_rgbx[:,:,:3] = latents_ubyte
self.event = torch.cuda.Event()
self.event.record()
preview_image = Image.frombuffer('RGBX', (latents_ubyte.shape[1], latents_ubyte.shape[0]),
latents_rgbx.numpy().data, 'raw', 'RGBX', 0, 1)
return ("JPEG", preview_image, MAX_PREVIEW_RESOLUTION, self.event)
latents_ubyte = (((preview_tensor + 1.0) / 2.0).clamp(0, 1) # change scale from -1..1 to 0..1
.mul(0xFF) # to 0..255
).to(device="cpu", dtype=torch.uint8, non_blocking=comfy.model_management.device_supports_non_blocking(preview_tensor.device))
latents_rgbx = torch.zeros(latents_ubyte.shape[:2] + (4,), device="cpu", dtype=torch.uint8)
latents_rgbx[:,:,:3] = latents_ubyte
self.event = torch.cuda.Event()
self.event.record()
preview_image = Image.frombuffer('RGBX', (latents_ubyte.shape[1], latents_ubyte.shape[0]),
latents_rgbx.numpy().data, 'raw', 'RGBX', 0, 1)
return ("JPEG", preview_image, MAX_PREVIEW_RESOLUTION, self.event)
).to(device="cpu", dtype=torch.uint8, non_blocking=False)
return ("JPEG", Image.fromarray(latents_ubyte.numpy()), MAX_PREVIEW_RESOLUTION)
class TAESDPreviewerImpl(LatentPreviewer):
def __init__(self, taesd):