jinx/ComfyUI
1
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2025-08-02 23:14:49 +08:00
Code Issues Activity

Make --gpu-only put intermediate values in GPU memory instead of cpu.

Browse Source
This commit is contained in:
comfyanonymous
2023-12-08 02:35:45 -05:00
parent cdff081023
commit 9ac0b487ac
9 changed files with 36 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
comfy/sd1_clip.py
View File

@@ -39,7 +39,7 @@ class ClipTokenWeightEncoder:
out, pooled = self.encode(to_encode)
if pooled is not None:
first_pooled = pooled[0:1].cpu()
first_pooled = pooled[0:1].to(model_management.intermediate_device())
else:
first_pooled = pooled
@@ -56,8 +56,8 @@ class ClipTokenWeightEncoder:
output.append(z)
if (len(output) == 0):
return out[-1:].cpu(), first_pooled
return torch.cat(output, dim=-2).cpu(), first_pooled
return out[-1:].to(model_management.intermediate_device()), first_pooled
return torch.cat(output, dim=-2).to(model_management.intermediate_device()), first_pooled
class SDClipModel(torch.nn.Module, ClipTokenWeightEncoder):
"""Uses the CLIP transformer encoder for text (from huggingface)"""