mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-12 09:27:15 +08:00
Compare commits
10 Commits
9f624af5da
...
2b2f46fa11
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2b2f46fa11 | ||
|
|
51bf508a0b | ||
|
|
a3020f107e | ||
|
|
7cf4e78335 | ||
|
|
7747c342d4 | ||
|
|
439bd807f8 | ||
|
|
b08debceca | ||
|
|
000c6b784e | ||
|
|
45a2363e6a | ||
|
|
edd44a6874 |
91
.github/workflows/cla.yml
vendored
Normal file
91
.github/workflows/cla.yml
vendored
Normal file
@ -0,0 +1,91 @@
|
||||
name: CLA Assistant
|
||||
|
||||
on:
|
||||
issue_comment:
|
||||
types: [created]
|
||||
pull_request_target:
|
||||
types: [opened, synchronize, closed]
|
||||
|
||||
permissions:
|
||||
actions: write
|
||||
contents: read # 'read' is enough because signatures live in a REMOTE repo
|
||||
pull-requests: write
|
||||
statuses: write
|
||||
|
||||
jobs:
|
||||
cla-assistant:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
# The CLA action normally requires every commit author in a PR to sign.
|
||||
# We only want the PR author to sign, so we allowlist all other committers
|
||||
# by computing them from the PR's commits and excluding the PR author.
|
||||
- name: Build author-only allowlist
|
||||
id: allowlist
|
||||
if: >
|
||||
github.event_name == 'pull_request_target' ||
|
||||
(github.event_name == 'issue_comment' && github.event.issue.pull_request && (
|
||||
github.event.comment.body == 'recheck' ||
|
||||
github.event.comment.body == 'I have read and agree to the Contributor License Agreement'
|
||||
))
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }}
|
||||
PR_AUTHOR: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
|
||||
BASE_ALLOWLIST: action@github.com,actions-user,ampagent,claude,comfy-pr-bot,GitHub Action,github-actions,github-actions[bot],Glary Bot,Glary-Bot,*[bot]
|
||||
run: |
|
||||
others=$(gh api "repos/${{ github.repository }}/pulls/${PR_NUMBER}/commits" --paginate \
|
||||
--jq '.[] | (.author.login // empty), (.committer.login // empty)' \
|
||||
| sort -u | grep -vix "${PR_AUTHOR}" | paste -sd, -)
|
||||
if [ -n "$others" ]; then
|
||||
echo "allowlist=${BASE_ALLOWLIST},${others}" >> "$GITHUB_OUTPUT"
|
||||
else
|
||||
echo "allowlist=${BASE_ALLOWLIST}" >> "$GITHUB_OUTPUT"
|
||||
fi
|
||||
|
||||
- name: CLA Assistant
|
||||
# Run on PR events, on "recheck" comment, or when someone posts the exact signing phrase.
|
||||
# IMPORTANT: this phrase must match `custom-pr-sign-comment` below.
|
||||
if: >
|
||||
github.event_name == 'pull_request_target' ||
|
||||
(github.event_name == 'issue_comment' && github.event.issue.pull_request && (
|
||||
github.event.comment.body == 'recheck' ||
|
||||
github.event.comment.body == 'I have read and agree to the Contributor License Agreement'
|
||||
))
|
||||
uses: contributor-assistant/github-action@ca4a40a7d1004f18d9960b404b97e5f30a505a08 # v2.6.1
|
||||
env:
|
||||
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
|
||||
# PAT required to write to the centralized signatures repo.
|
||||
PERSONAL_ACCESS_TOKEN: ${{ secrets.PERSONAL_ACCESS_TOKEN }}
|
||||
with:
|
||||
# Where the CLA document lives (shown to contributors)
|
||||
path-to-document: https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md
|
||||
|
||||
# Centralized signature storage
|
||||
remote-organization-name: comfy-org
|
||||
remote-repository-name: comfy-cla
|
||||
path-to-signatures: signatures/cla.json
|
||||
branch: main
|
||||
|
||||
# Only the PR author must sign: bots plus every non-author committer
|
||||
# are allowlisted via the "Build author-only allowlist" step above.
|
||||
# *[bot] is a catch-all for any GitHub App bot account.
|
||||
allowlist: ${{ steps.allowlist.outputs.allowlist }}
|
||||
|
||||
# Custom PR comment messages
|
||||
custom-notsigned-prcomment: |
|
||||
🎉 Thank you for your contribution, we really appreciate it! 🎉
|
||||
|
||||
Like many open source projects, we require contributors to sign our [Contributor License Agreement (CLA)](https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md). A CLA makes the ownership of contributions explicit, so contributors and the project share a clear understanding of how the code can be used. By signing, you:
|
||||
|
||||
- Confirm that you own your contribution.
|
||||
- Keep the right to reuse your own code.
|
||||
- Grant us a copyright license to include and share it within our projects.
|
||||
|
||||
CLAs are standard practice across major open source projects including those under the Apache Software Foundation and the Linux Foundation. Ours is based on the Apache Software Foundation's CLA. Most importantly, it would enable us to relicense the project under a more permissive license in the future, giving the project and its community greater flexibility.
|
||||
|
||||
✍ **To sign, please post a new comment on this PR with exactly the following text:** ✍
|
||||
|
||||
custom-pr-sign-comment: I have read and agree to the Contributor License Agreement
|
||||
|
||||
custom-allsigned-prcomment: |
|
||||
✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged.
|
||||
@ -70,6 +70,11 @@ def stochastic_rounding(value, dtype, seed=0):
|
||||
if dtype == torch.bfloat16:
|
||||
return value.to(dtype=torch.bfloat16)
|
||||
if dtype == torch.float8_e4m3fn or dtype == torch.float8_e5m2:
|
||||
# MPS does not support FP8 dtypes — perform rounding on CPU and return the result there.
|
||||
on_mps = value.device.type == "mps"
|
||||
if on_mps:
|
||||
value = value.cpu()
|
||||
|
||||
generator = torch.Generator(device=value.device)
|
||||
generator.manual_seed(seed)
|
||||
if _CK_STOCHASTIC_ROUNDING_AVAILABLE:
|
||||
@ -178,6 +183,12 @@ def stochastic_round_quantize_nvfp4(x, per_tensor_scale, pad_16x, seed=0):
|
||||
"""Round up x to the nearest multiple."""
|
||||
return ((x + multiple - 1) // multiple) * multiple
|
||||
|
||||
# MPS does not support FP8 dtypes used for block scales — perform on CPU.
|
||||
on_mps = x.device.type == "mps"
|
||||
if on_mps:
|
||||
x = x.cpu()
|
||||
per_tensor_scale = per_tensor_scale.cpu() if isinstance(per_tensor_scale, torch.Tensor) else per_tensor_scale
|
||||
|
||||
generator = torch.Generator(device=x.device)
|
||||
generator.manual_seed(seed)
|
||||
|
||||
@ -198,6 +209,12 @@ def stochastic_round_quantize_nvfp4_by_block(x, per_tensor_scale, pad_16x, seed=
|
||||
"""Round up x to the nearest multiple."""
|
||||
return ((x + multiple - 1) // multiple) * multiple
|
||||
|
||||
# MPS does not support FP8 dtypes used for block scales — perform on CPU.
|
||||
on_mps = x.device.type == "mps"
|
||||
if on_mps:
|
||||
x = x.cpu()
|
||||
per_tensor_scale = per_tensor_scale.cpu() if isinstance(per_tensor_scale, torch.Tensor) else per_tensor_scale
|
||||
|
||||
orig_shape = x.shape
|
||||
|
||||
# Handle padding
|
||||
|
||||
@ -97,6 +97,12 @@ class _TensorCoreFP8LayoutBase(_CKFp8Layout):
|
||||
if not isinstance(scale, torch.Tensor):
|
||||
scale = torch.tensor(scale, device=tensor.device, dtype=torch.float32)
|
||||
|
||||
# MPS does not support FP8 dtypes — move to CPU for quantization.
|
||||
on_mps = tensor.device.type == "mps"
|
||||
if on_mps:
|
||||
tensor = tensor.cpu()
|
||||
scale = scale.cpu()
|
||||
|
||||
if stochastic_rounding > 0:
|
||||
if inplace_ops:
|
||||
tensor *= (1.0 / scale).to(tensor.dtype)
|
||||
|
||||
@ -468,6 +468,9 @@ class CLIP:
|
||||
def decode(self, token_ids, skip_special_tokens=True):
|
||||
return self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
|
||||
|
||||
def is_dynamic(self):
|
||||
return self.patcher.is_dynamic()
|
||||
|
||||
class VAE:
|
||||
def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None):
|
||||
if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
|
||||
@ -1251,6 +1254,8 @@ class VAE:
|
||||
except:
|
||||
return None
|
||||
|
||||
def is_dynamic(self):
|
||||
return self.patcher.is_dynamic()
|
||||
|
||||
class StyleModel:
|
||||
def __init__(self, model, device="cpu"):
|
||||
|
||||
@ -937,22 +937,41 @@ class BaseGenerate:
|
||||
return torch.argmax(logits, dim=-1, keepdim=True)
|
||||
|
||||
# Sampling mode
|
||||
if repetition_penalty != 1.0:
|
||||
for i in range(logits.shape[0]):
|
||||
for token_id in set(token_history):
|
||||
logits[i, token_id] *= repetition_penalty if logits[i, token_id] < 0 else 1/repetition_penalty
|
||||
|
||||
if presence_penalty is not None and presence_penalty != 0.0:
|
||||
for i in range(logits.shape[0]):
|
||||
for token_id in set(token_history):
|
||||
logits[i, token_id] -= presence_penalty
|
||||
if len(token_history) > 0 and (repetition_penalty != 1.0 or (presence_penalty is not None and presence_penalty != 0.0)):
|
||||
token_ids = torch.tensor(list(set(token_history)), device=logits.device)
|
||||
token_logits = logits[:, token_ids]
|
||||
if repetition_penalty != 1.0:
|
||||
token_logits = torch.where(token_logits < 0, token_logits * repetition_penalty, token_logits / repetition_penalty)
|
||||
if presence_penalty is not None and presence_penalty != 0.0:
|
||||
token_logits = token_logits - presence_penalty
|
||||
logits[:, token_ids] = token_logits
|
||||
|
||||
if temperature != 1.0:
|
||||
logits = logits / temperature
|
||||
|
||||
if top_k > 0:
|
||||
indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
|
||||
logits[indices_to_remove] = torch.finfo(logits.dtype).min
|
||||
top_k = min(top_k, logits.shape[-1])
|
||||
logits, top_indices = torch.topk(logits, top_k)
|
||||
|
||||
if min_p > 0.0:
|
||||
probs_before_filter = torch.nn.functional.softmax(logits, dim=-1)
|
||||
top_probs, _ = probs_before_filter.max(dim=-1, keepdim=True)
|
||||
min_threshold = min_p * top_probs
|
||||
indices_to_remove = probs_before_filter < min_threshold
|
||||
logits[indices_to_remove] = torch.finfo(logits.dtype).min
|
||||
|
||||
if top_p < 1.0:
|
||||
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
|
||||
cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1)
|
||||
sorted_indices_to_remove = cumulative_probs > top_p
|
||||
sorted_indices_to_remove[..., 0] = False
|
||||
indices_to_remove = torch.zeros_like(logits, dtype=torch.bool)
|
||||
indices_to_remove.scatter_(1, sorted_indices, sorted_indices_to_remove)
|
||||
logits[indices_to_remove] = torch.finfo(logits.dtype).min
|
||||
|
||||
probs = torch.nn.functional.softmax(logits, dim=-1)
|
||||
next_token = torch.multinomial(probs, num_samples=1, generator=generator)
|
||||
return top_indices.gather(1, next_token)
|
||||
|
||||
if min_p > 0.0:
|
||||
probs_before_filter = torch.nn.functional.softmax(logits, dim=-1)
|
||||
|
||||
@ -281,11 +281,18 @@ class VideoFromFile(VideoInput):
|
||||
video_done = False
|
||||
audio_done = True
|
||||
|
||||
if len(container.streams.audio):
|
||||
audio_stream = container.streams.audio[-1]
|
||||
# Use the last decodable audio stream. Streams FFmpeg has no decoder for have no codec context,
|
||||
# and decoding their packets crashes the process. (e.g. APAC spatial-audio track in iPhone)
|
||||
audio_stream = next(
|
||||
(s for s in reversed(container.streams.audio) if s.codec_context is not None),
|
||||
None,
|
||||
)
|
||||
if audio_stream is not None:
|
||||
streams += [audio_stream]
|
||||
resampler = av.audio.resampler.AudioResampler(format='fltp')
|
||||
audio_done = False
|
||||
elif len(container.streams.audio):
|
||||
logging.warning("No decodable audio stream found in video; ignoring audio.")
|
||||
|
||||
for packet in container.demux(*streams):
|
||||
if video_done and audio_done:
|
||||
@ -457,10 +464,13 @@ class VideoFromFile(VideoInput):
|
||||
else:
|
||||
output_container.metadata[key] = json.dumps(value)
|
||||
|
||||
# Add streams to the new container
|
||||
# Add streams to the new container. Streams with no codec context cannot be used as an output template.
|
||||
stream_map = {}
|
||||
for stream in streams:
|
||||
if isinstance(stream, (av.VideoStream, av.AudioStream, SubtitleStream)):
|
||||
if stream.codec_context is None:
|
||||
logging.warning("Skipping %s stream %d with unsupported codec", stream.type, stream.index)
|
||||
continue
|
||||
out_stream = output_container.add_stream_from_template(template=stream, opaque=True)
|
||||
stream_map[stream] = out_stream
|
||||
|
||||
|
||||
@ -158,7 +158,14 @@ async def upload_video_to_comfyapi(
|
||||
|
||||
# Convert VideoInput to BytesIO using specified container/codec
|
||||
video_bytes_io = BytesIO()
|
||||
video.save_to(video_bytes_io, format=container, codec=codec)
|
||||
try:
|
||||
video.save_to(video_bytes_io, format=container, codec=codec)
|
||||
except Exception as e:
|
||||
raise ValueError(
|
||||
f"Could not convert the input video to {container.value.upper()} for upload; "
|
||||
f"the file may be corrupted or use an unsupported codec. "
|
||||
f"Try re-exporting it as MP4 (H.264). Original error: {e}"
|
||||
) from e
|
||||
video_bytes_io.seek(0)
|
||||
|
||||
return await upload_file_to_comfyapi(cls, video_bytes_io, filename, upload_mime_type, wait_label)
|
||||
|
||||
@ -503,6 +503,21 @@ RAM_CACHE_DEFAULT_RAM_USAGE = 0.05
|
||||
|
||||
RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3
|
||||
|
||||
|
||||
def all_outputs_dynamic(outputs):
|
||||
if outputs is None:
|
||||
return False
|
||||
|
||||
for output in outputs:
|
||||
if isinstance(output, (list, tuple)):
|
||||
if not all_outputs_dynamic(output):
|
||||
return False
|
||||
elif not hasattr(output, "is_dynamic") or not output.is_dynamic():
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
|
||||
class RAMPressureCache(LRUCache):
|
||||
|
||||
def __init__(self, key_class, enable_providers=False):
|
||||
@ -533,7 +548,11 @@ class RAMPressureCache(LRUCache):
|
||||
for key, cache_entry in self.cache.items():
|
||||
if not free_active and self.used_generation[key] == self.generation:
|
||||
continue
|
||||
oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
|
||||
|
||||
if all_outputs_dynamic(cache_entry.outputs) and self.used_generation[key] == self.generation:
|
||||
continue
|
||||
|
||||
oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
|
||||
|
||||
ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE
|
||||
def scan_list_for_ram_usage(outputs):
|
||||
|
||||
150
comfy_extras/nodes_text_overlay.py
Normal file
150
comfy_extras/nodes_text_overlay.py
Normal file
@ -0,0 +1,150 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image as PILImage, ImageColor, ImageDraw, ImageFont
|
||||
from typing_extensions import override
|
||||
|
||||
from comfy_api.latest import ComfyExtension, IO
|
||||
|
||||
|
||||
class TextOverlay(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="TextOverlay",
|
||||
display_name="Draw Text Overlay",
|
||||
category="text",
|
||||
description="Draw text overlay on an image or batch of images.",
|
||||
search_aliases=["text", "label", "caption", "subtitle", "watermark", "title", "addlabel", "overlay"],
|
||||
inputs=[
|
||||
IO.Image.Input("images"),
|
||||
IO.String.Input("text", multiline=True, default=""),
|
||||
IO.Float.Input("font_size", default=5.0, min=0.5, max=50.0, step=0.5, tooltip="Font size as a percentage of the image height."),
|
||||
IO.Color.Input("color", default="#ffffff", tooltip="Color of the text."),
|
||||
IO.Combo.Input("position", options=["top", "bottom"], default="top"),
|
||||
IO.Combo.Input("align", options=["left", "center", "right"], default="left"),
|
||||
IO.Boolean.Input("outline", default=True, tooltip="Draw a black outline around the text."),
|
||||
],
|
||||
outputs=[IO.Image.Output(display_name="images")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images, text, font_size, color, position, align, outline) -> IO.NodeOutput:
|
||||
if text.strip() == "":
|
||||
return IO.NodeOutput(images)
|
||||
|
||||
text = text.replace("\\n", "\n").replace("\\t", "\t")
|
||||
|
||||
text_rgba = cls.parse_color_to_rgba(color)
|
||||
outline_rgba = (0, 0, 0, 255) if outline else (0, 0, 0, 0)
|
||||
|
||||
# Render the overlay once and composite it across all frames in the batch
|
||||
height = images.shape[1]
|
||||
width = images.shape[2]
|
||||
overlay_rgb, overlay_alpha = cls.render_overlay_text(width, height, text, position, align, font_size, text_rgba, outline_rgba)
|
||||
overlay_rgb = overlay_rgb.to(device=images.device, dtype=images.dtype)
|
||||
overlay_alpha = overlay_alpha.to(device=images.device, dtype=images.dtype)
|
||||
|
||||
result = images * (1.0 - overlay_alpha) + overlay_rgb * overlay_alpha
|
||||
return IO.NodeOutput(result)
|
||||
|
||||
@staticmethod
|
||||
def parse_color_to_rgba(color_string):
|
||||
parsed = ImageColor.getrgb(color_string)
|
||||
|
||||
if len(parsed) == 3:
|
||||
return (*parsed, 255)
|
||||
|
||||
return parsed
|
||||
|
||||
@classmethod
|
||||
def render_overlay_text(cls, width, height, text, position, align, font_size, text_rgba, outline_rgba):
|
||||
line_spacing = 1.2
|
||||
margin_percent = 1.0
|
||||
min_font_percent = 2.0
|
||||
min_font_pixels = 10
|
||||
outline_thickness_factor = 0.04
|
||||
|
||||
# Draw onto a transparent layer so the result can be alpha-composited over any frame.
|
||||
layer = PILImage.new("RGBA", (width, height), (0, 0, 0, 0))
|
||||
draw = ImageDraw.Draw(layer)
|
||||
|
||||
margin = int(round(margin_percent / 100.0 * min(width, height)))
|
||||
max_width = max(1, width - 2 * margin)
|
||||
max_height = max(1, height - 2 * margin)
|
||||
|
||||
# Font scales with resolution, then shrinks to fit the height.
|
||||
size = max(1, int(round(font_size / 100.0 * height)))
|
||||
floor = min(size, max(min_font_pixels, int(round(min_font_percent / 100.0 * height))))
|
||||
|
||||
while True:
|
||||
font = ImageFont.load_default(size=size)
|
||||
stroke = max(1, int(round(size * outline_thickness_factor))) if outline_rgba[3] > 0 else 0
|
||||
block = "\n".join(cls.wrap_text(text, font, max_width))
|
||||
# convert line spacing to pixel spacing
|
||||
single = draw.textbbox((0, 0), "Ay", font=font, stroke_width=stroke)
|
||||
double = draw.multiline_textbbox((0, 0), "Ay\nAy", font=font, spacing=0, stroke_width=stroke)
|
||||
natural_advance = (double[3] - double[1]) - (single[3] - single[1])
|
||||
pixel_spacing = int(round(size * line_spacing - natural_advance))
|
||||
box = draw.multiline_textbbox((0, 0), block, font=font, spacing=pixel_spacing, stroke_width=stroke)
|
||||
block_height = box[3] - box[1]
|
||||
|
||||
if block_height <= max_height or size <= floor:
|
||||
break
|
||||
|
||||
size = max(floor, int(size * 0.9))
|
||||
|
||||
anchor_h, x = {"left": ("l", margin), "center": ("m", width / 2), "right": ("r", width - margin)}[align]
|
||||
|
||||
# Offset y so the rendered text sits flush against the margin
|
||||
if position == "bottom":
|
||||
y = height - margin - box[3]
|
||||
else:
|
||||
y = margin - box[1]
|
||||
|
||||
draw.multiline_text((x, y), block, font=font, fill=text_rgba, anchor=anchor_h + "a",
|
||||
align=align, spacing=pixel_spacing, stroke_width=stroke, stroke_fill=outline_rgba)
|
||||
|
||||
overlay = np.array(layer).astype(np.float32) / 255.0
|
||||
overlay_rgb = torch.from_numpy(overlay[:, :, :3])
|
||||
overlay_alpha = torch.from_numpy(overlay[:, :, 3:4])
|
||||
return overlay_rgb, overlay_alpha
|
||||
|
||||
@staticmethod
|
||||
def wrap_text(text, font, max_width):
|
||||
lines = []
|
||||
for raw_line in text.split("\n"):
|
||||
words = raw_line.split()
|
||||
if not words:
|
||||
lines.append("")
|
||||
continue
|
||||
current = ""
|
||||
# Break the line into words and split words that are too long
|
||||
for word in words:
|
||||
while font.getlength(word) > max_width and len(word) > 1:
|
||||
cut = 1
|
||||
while cut < len(word) and font.getlength(word[:cut + 1]) <= max_width:
|
||||
cut += 1
|
||||
if current:
|
||||
lines.append(current)
|
||||
current = ""
|
||||
lines.append(word[:cut])
|
||||
word = word[cut:]
|
||||
candidate = word if not current else current + " " + word
|
||||
if not current or font.getlength(candidate) <= max_width:
|
||||
current = candidate
|
||||
else:
|
||||
lines.append(current)
|
||||
current = word
|
||||
if current:
|
||||
lines.append(current)
|
||||
return lines
|
||||
|
||||
|
||||
class TextOverlayExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [TextOverlay]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> TextOverlayExtension:
|
||||
return TextOverlayExtension()
|
||||
1
nodes.py
1
nodes.py
@ -2478,6 +2478,7 @@ async def init_builtin_extra_nodes():
|
||||
"nodes_glsl.py",
|
||||
"nodes_lora_debug.py",
|
||||
"nodes_textgen.py",
|
||||
"nodes_text_overlay.py",
|
||||
"nodes_color.py",
|
||||
"nodes_toolkit.py",
|
||||
"nodes_replacements.py",
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
comfyui-frontend-package==1.45.20
|
||||
comfyui-workflow-templates==0.11.2
|
||||
comfyui-embedded-docs==0.5.6
|
||||
comfyui-embedded-docs==0.5.7
|
||||
torch
|
||||
torchsde
|
||||
torchvision
|
||||
|
||||
147
tests/test_fp8_mps.py
Normal file
147
tests/test_fp8_mps.py
Normal file
@ -0,0 +1,147 @@
|
||||
"""
|
||||
Tests for FP8 quantization on MPS (Apple Silicon) devices.
|
||||
|
||||
MPS does not natively support float8_e4m3fn or float8_e5m2 dtypes.
|
||||
These tests verify that:
|
||||
1. FP8 operations correctly fall back to CPU when on MPS.
|
||||
2. The round-trip (quantize on CPU -> result on original device) is numerically sound.
|
||||
3. No "Placeholder storage has not been allocated on MPS device!" errors occur.
|
||||
"""
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import comfy.float
|
||||
from comfy.quant_ops import TensorCoreFP8E4M3Layout, TensorCoreFP8E5M2Layout
|
||||
|
||||
# Skip the entire module if MPS is not available
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not torch.backends.mps.is_available(),
|
||||
reason="MPS backend not available"
|
||||
)
|
||||
|
||||
# ── helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
def _make_mps_tensor(shape=(256, 256), dtype=torch.float32):
|
||||
return torch.randn(shape, device="mps", dtype=dtype)
|
||||
|
||||
|
||||
# ── Tests for comfy.float ────────────────────────────────────────────────────
|
||||
|
||||
class TestStochasticRoundingMPS:
|
||||
"""Tests for comfy.float.stochastic_rounding on MPS device."""
|
||||
|
||||
def test_stochastic_rounding_fp8_e4m3fn_on_mps(self):
|
||||
"""stochastic_rounding must not crash when input is on MPS and target dtype is float8_e4m3fn."""
|
||||
x = _make_mps_tensor((64, 64), dtype=torch.float32)
|
||||
result = comfy.float.stochastic_rounding(x, dtype=torch.float8_e4m3fn, seed=42)
|
||||
|
||||
assert result.dtype == torch.float8_e4m3fn
|
||||
assert result.shape == x.shape
|
||||
|
||||
def test_stochastic_rounding_fp8_e5m2_on_mps(self):
|
||||
"""stochastic_rounding must not crash when input is on MPS and target dtype is float8_e5m2."""
|
||||
x = _make_mps_tensor((64, 64), dtype=torch.float32)
|
||||
result = comfy.float.stochastic_rounding(x, dtype=torch.float8_e5m2, seed=42)
|
||||
|
||||
assert result.dtype == torch.float8_e5m2
|
||||
assert result.shape == x.shape
|
||||
|
||||
def test_stochastic_rounding_fp8_result_on_cpu(self):
|
||||
"""Result of FP8 rounding from MPS input should be on CPU (since MPS can't hold FP8)."""
|
||||
x = _make_mps_tensor((32, 32), dtype=torch.float32)
|
||||
result = comfy.float.stochastic_rounding(x, dtype=torch.float8_e4m3fn, seed=42)
|
||||
|
||||
# FP8 tensors cannot live on MPS, so result must be on CPU
|
||||
assert result.device.type == "cpu"
|
||||
|
||||
def test_stochastic_rounding_non_fp8_still_works(self):
|
||||
"""Non-FP8 dtypes on MPS must still work as before (no regression)."""
|
||||
x = _make_mps_tensor((32, 32), dtype=torch.float32)
|
||||
|
||||
r16 = comfy.float.stochastic_rounding(x, dtype=torch.float16, seed=0)
|
||||
assert r16.dtype == torch.float16
|
||||
assert r16.device.type == "mps"
|
||||
|
||||
rbf16 = comfy.float.stochastic_rounding(x, dtype=torch.bfloat16, seed=0)
|
||||
assert rbf16.dtype == torch.bfloat16
|
||||
assert rbf16.device.type == "mps"
|
||||
|
||||
def test_stochastic_rounding_fp8_numerical_sanity(self):
|
||||
"""FP8 round-trip (float32 -> fp8 -> float32) should have bounded error."""
|
||||
x = torch.randn(128, 128, device="mps", dtype=torch.float32)
|
||||
x_clamped = torch.clamp(x, min=-448, max=448) # FP8 e4m3fn range
|
||||
|
||||
fp8 = comfy.float.stochastic_rounding(x_clamped, dtype=torch.float8_e4m3fn, seed=123)
|
||||
# Convert back to float32 for comparison
|
||||
reconstructed = fp8.to(torch.float32)
|
||||
|
||||
# Max relative error should be bounded (FP8 e4m3fn has ~0.125 relative precision)
|
||||
x_cpu = x_clamped.cpu()
|
||||
max_abs_err = (reconstructed - x_cpu).abs().max().item()
|
||||
# FP8 e4m3fn max value is 448, min subnormal ~0.001953
|
||||
# For random normal data, error should be well under 1.0
|
||||
assert max_abs_err < 2.0, f"FP8 round-trip error too large: {max_abs_err}"
|
||||
|
||||
|
||||
class TestManualStochasticRoundMPS:
|
||||
"""Tests for comfy.float.manual_stochastic_round_to_float8 on MPS device."""
|
||||
|
||||
def test_manual_round_fp8_on_mps_tensor(self):
|
||||
"""stochastic_rounding handles MPS generator internally without 'Placeholder storage' error."""
|
||||
x = _make_mps_tensor((16, 16), dtype=torch.float32)
|
||||
result = comfy.float.stochastic_rounding(x, dtype=torch.float8_e4m3fn, seed=42)
|
||||
assert result.dtype == torch.float8_e4m3fn
|
||||
|
||||
|
||||
class TestNVFP4StochasticRoundMPS:
|
||||
"""Tests for NVFP4 stochastic rounding on MPS - also creates FP8 tensors internally."""
|
||||
|
||||
def test_nvfp4_stochastic_round_on_mps(self):
|
||||
"""stochastic_round_quantize_nvfp4 creates FP8 block scales internally."""
|
||||
# NVFP4 requires 2D input with dimensions divisible by 16
|
||||
x = torch.randn(32, 32, device="mps", dtype=torch.float32)
|
||||
scale = torch.tensor(1.0, device="mps", dtype=torch.float32)
|
||||
|
||||
# This should not crash - internally creates float8_e4m3fn block scales
|
||||
qdata, block_scale = comfy.float.stochastic_round_quantize_nvfp4(
|
||||
x, scale, pad_16x=False, seed=42
|
||||
)
|
||||
assert qdata.dtype == torch.uint8
|
||||
|
||||
|
||||
# ── Tests for comfy.quant_ops (integration) ──────────────────────────────────
|
||||
|
||||
class TestQuantOpsMPS:
|
||||
"""Tests for the quantization ops layer that calls into comfy.float."""
|
||||
|
||||
def test_fp8_layout_quantize_on_mps(self):
|
||||
"""TensorCoreFP8E4M3Layout.quantize must work with MPS tensors."""
|
||||
x = _make_mps_tensor((64, 64), dtype=torch.bfloat16)
|
||||
qdata, params = TensorCoreFP8E4M3Layout.quantize(
|
||||
x, scale="recalculate", stochastic_rounding=42
|
||||
)
|
||||
|
||||
assert qdata.dtype == torch.float8_e4m3fn
|
||||
assert params.orig_dtype == torch.bfloat16
|
||||
|
||||
def test_fp8_layout_quantize_without_stochastic_on_mps(self):
|
||||
"""TensorCoreFP8E4M3Layout.quantize with stochastic_rounding=0 uses ck.quantize_per_tensor_fp8."""
|
||||
x = _make_mps_tensor((64, 64), dtype=torch.bfloat16)
|
||||
qdata, params = TensorCoreFP8E4M3Layout.quantize(
|
||||
x, scale="recalculate", stochastic_rounding=0
|
||||
)
|
||||
|
||||
assert qdata.dtype == torch.float8_e4m3fn
|
||||
|
||||
def test_fp8_e5m2_layout_quantize_on_mps(self):
|
||||
"""TensorCoreFP8E5M2Layout.quantize must work with MPS tensors."""
|
||||
x = _make_mps_tensor((64, 64), dtype=torch.float32)
|
||||
qdata, params = TensorCoreFP8E5M2Layout.quantize(
|
||||
x, scale="recalculate", stochastic_rounding=42
|
||||
)
|
||||
|
||||
assert qdata.dtype == torch.float8_e5m2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "--tb=short"])
|
||||
Loading…
Reference in New Issue
Block a user