Merge branch 'master' into master

This commit is contained in:
Li Xiang 2025-12-30 16:58:28 +08:00 committed by GitHub
commit 68f852aff5
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 29 additions and 10 deletions

View File

@ -1126,6 +1126,16 @@ if not args.disable_pinned_memory:
PINNING_ALLOWED_TYPES = set(["Parameter", "QuantizedTensor"]) PINNING_ALLOWED_TYPES = set(["Parameter", "QuantizedTensor"])
def discard_cuda_async_error():
try:
a = torch.tensor([1], dtype=torch.uint8, device=get_torch_device())
b = torch.tensor([1], dtype=torch.uint8, device=get_torch_device())
_ = a + b
torch.cuda.synchronize()
except torch.AcceleratorError:
#Dump it! We already know about it from the synchronous return
pass
def pin_memory(tensor): def pin_memory(tensor):
global TOTAL_PINNED_MEMORY global TOTAL_PINNED_MEMORY
if MAX_PINNED_MEMORY <= 0: if MAX_PINNED_MEMORY <= 0:
@ -1158,6 +1168,9 @@ def pin_memory(tensor):
PINNED_MEMORY[ptr] = size PINNED_MEMORY[ptr] = size
TOTAL_PINNED_MEMORY += size TOTAL_PINNED_MEMORY += size
return True return True
else:
logging.warning("Pin error.")
discard_cuda_async_error()
return False return False
@ -1186,6 +1199,9 @@ def unpin_memory(tensor):
if len(PINNED_MEMORY) == 0: if len(PINNED_MEMORY) == 0:
TOTAL_PINNED_MEMORY = 0 TOTAL_PINNED_MEMORY = 0
return True return True
else:
logging.warning("Unpin error.")
discard_cuda_async_error()
return False return False

View File

@ -667,16 +667,19 @@ class ResizeImagesByLongerEdgeNode(ImageProcessingNode):
@classmethod @classmethod
def _process(cls, image, longer_edge): def _process(cls, image, longer_edge):
img = tensor_to_pil(image) resized_images = []
w, h = img.size for image_i in image:
if w > h: img = tensor_to_pil(image_i)
new_w = longer_edge w, h = img.size
new_h = int(h * (longer_edge / w)) if w > h:
else: new_w = longer_edge
new_h = longer_edge new_h = int(h * (longer_edge / w))
new_w = int(w * (longer_edge / h)) else:
img = img.resize((new_w, new_h), Image.Resampling.LANCZOS) new_h = longer_edge
return pil_to_tensor(img) new_w = int(w * (longer_edge / h))
img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)
resized_images.append(pil_to_tensor(img))
return torch.cat(resized_images, dim=0)
class CenterCropImagesNode(ImageProcessingNode): class CenterCropImagesNode(ImageProcessingNode):