From 32e7557facaaaddab71d29da9ad6beed047bfe6e Mon Sep 17 00:00:00 2001 From: mcDandy Date: Sun, 18 Oct 2026 01:36:00 +0200 Subject: [PATCH] AI: audit 2 --- more_math/Parser/UnifiedMathVisitor.py | 762 ++++++++-------------- more_math/Parser/antlr_router.py | 18 +- more_math/Parser/optical_flow_utils.py | 136 +--- more_math/SelectiveGuiderMathNode.py | 25 +- more_math/helper_functions.py | 34 +- more_math/modelLikeCommon.py | 31 +- tests/test_unified_visitor_regressions.py | 48 +- 7 files changed, 357 insertions(+), 697 deletions(-) diff --git a/more_math/Parser/UnifiedMathVisitor.py b/more_math/Parser/UnifiedMathVisitor.py index 6f39949..8fd0469 100644 --- a/more_math/Parser/UnifiedMathVisitor.py +++ b/more_math/Parser/UnifiedMathVisitor.py @@ -1,5 +1,8 @@ +import bisect +import functools import time import os +import sys import torch import math import inspect @@ -22,15 +25,17 @@ from PIL import Image, ImageDraw, ImageFont import numpy as np -class ReturnSignal: - __slots__ = ("value",) +class ControlSignal(Exception): + """Returned by break/continue/return visitors; visit() throws it into the waiting parent visitors.""" + +class ReturnSignal(ControlSignal): def __init__(self, value): self.value = value -class BreakSignal: +class BreakSignal(ControlSignal): pass -class ContinueSignal: +class ContinueSignal(ControlSignal): pass class MathDict(dict): @@ -129,6 +134,140 @@ class ChangeTrackingDict(dict): self._changed_keys = set() +_FONT_WEIGHT_TIERS = ("thin", "extralight", "light", "regular", "medium", "semibold", "bold", "extrabold", "black") +_FONT_STYLE_SUFFIXES = { + "thin": ("thin", "hairline", "100"), + "extralight": ("extralight", "ultralight", "200"), + "light": ("light", "300"), + "regular": ("regular", "normal", "reg", "400", "standard"), + "medium": ("medium", "500"), + "semibold": ("semibold", "demibold", "600"), + "bold": ("bold", "bd", "700"), + "extrabold": ("extrabold", "ultrabold", "800"), + "black": ("black", "heavy", "blk", "900"), +} +# Compound names first so "semibold" is not read as "bold" or "extralight" as "light". +_FONT_WEIGHT_HINTS = tuple( + ((_FONT_WEIGHT_TIERS.index(tier) + 1) * 100, _FONT_STYLE_SUFFIXES[tier]) + for tier in ("extrabold", "semibold", "extralight", "black", "bold", "medium", "light", "thin") +) + + +def _normalise_font_name(value): + return value.lower().replace(" ", "").replace("-", "").replace("_", "") + + +def _font_axis_name(axis): + name = axis.get("name", "") + if isinstance(name, bytes): + name = name.decode(errors="ignore") + return str(name).lower() + + +@functools.lru_cache(maxsize=1) +def _system_fonts(): + """Index of installed fonts, scanned once per process (fonts installed later need a restart).""" + if sys.platform == "win32": + search_dirs = [os.path.join(os.environ.get("WINDIR", "C:\\Windows"), "Fonts"), os.path.expandvars(r"%LOCALAPPDATA%\Microsoft\Windows\Fonts")] + elif sys.platform == "darwin": + search_dirs = ["/Library/Fonts", "/System/Library/Fonts", os.path.expanduser("~/Library/Fonts")] + else: + search_dirs = [ + "/usr/share/fonts", + "/usr/local/share/fonts", + os.path.join(os.path.expanduser(os.environ.get("XDG_DATA_HOME", "~/.local/share")), "fonts"), + os.path.expanduser("~/.local/share/fonts"), + os.path.expanduser("~/.fonts"), + ] + + fonts = [] + for directory in dict.fromkeys(search_dirs): + for root, _, files in os.walk(directory): + for fname in files: + if not fname.lower().endswith((".ttf", ".otf", ".ttc", ".otc")): + continue + path = os.path.join(root, fname) + try: + fnt = ImageFont.truetype(path, 10) + except OSError: + continue + family, sub = fnt.getname() + try: + axes = {_font_axis_name(axis) for axis in fnt.get_variation_axes()} + except OSError: # not a variable font + axes = set() + fonts.append(( + path, + _normalise_font_name(os.path.splitext(fname)[0]), + _normalise_font_name(family or ""), + _normalise_font_name(sub or ""), + bool(axes & {"weight", "wght"}), + bool(axes & {"italic", "ital", "slant", "slnt"}), + )) + return fonts + + +def _find_font(name, tier, weight, italic): + """Path of the installed font best matching name/weight/italic, or None.""" + name_clean = _normalise_font_name(name) + suffixes = _FONT_STYLE_SUFFIXES[tier] + if italic: + suffixes = tuple(f"{s}{i}" for s in suffixes for i in ("italic", "ital", "it", "oblique")) + ("italic", "ital", "oblique", "it") + plain_regular = tier == "regular" and not italic + + best = None + for path, base, family, sub, weight_axis, italic_axis in _system_fonts(): + match = name_clean in base and ((plain_regular and base == name_clean) or any(s in base for s in suffixes)) + if not match and family and (name_clean in family or family in name_clean): + if (italic and italic_axis) or (not italic and weight_axis): + match = True + elif not (weight_axis or italic_axis): + if plain_regular: + match = sub == "" or any(s in sub for s in ("regular", "normal", "standard")) + else: + match = any(s in sub or s in family for s in suffixes) + if not match: + continue + + combined = f"{base} {sub}" + font_italic = any(t in combined for t in ("italic", "ital", "oblique")) or base.endswith((f"{name_clean}i", f"{name_clean}bi")) + font_weight = next((w for w, hints in _FONT_WEIGHT_HINTS if any(h in combined for h in hints)), 400) + if base.endswith((f"{name_clean}bd", f"{name_clean}bi")): + font_weight = 700 + + score = (0 if font_italic == italic else 10000) + (0 if weight_axis else abs(font_weight - weight)) + (0 if base == name_clean else 50) + if best is None or score < best[0]: + best = (score, path) + return best[1] if best is not None else None + + +def _apply_font_variations(fnt, weight, italic): + try: + axes = fnt.get_variation_axes() + except OSError: # not a variable font + return fnt + + values = [] + changed = False + for axis in axes: + axis_name = _font_axis_name(axis) + if axis_name in ("weight", "wght"): + target = weight + elif axis_name in ("italic", "ital"): + target = 1 if italic else 0 + elif axis_name in ("slant", "slnt"): + target = axis.get("minimum", -12) if italic else 0 + else: + values.append(axis.get("default", axis.get("minimum", 0))) + continue + values.append(max(axis["minimum"], min(axis["maximum"], target))) + changed = True + + if changed: + fnt.set_variation_by_axes(values) + return fnt + + class UnifiedMathVisitor(MathExprVisitor): def __init__(self, variables, shape=None, device=None, functions=None, depth=0, state_storage=None): @@ -153,46 +292,43 @@ class UnifiedMathVisitor(MathExprVisitor): if not inspect.isgenerator(gen): return gen + # Visitors yield child nodes instead of recursing, so deep scripts don't hit Python's recursion limit. + # A ControlSignal is thrown into the waiting parent; loops and calls catch the ones they handle. stack = [gen] last_result = None + signal = None while stack: try: - # If we're bubbling a signal, we need to check if the parent can handle it. - if isinstance(last_result, (ReturnSignal, BreakSignal, ContinueSignal)): - parent_gen = stack[-1] - func_name = parent_gen.gi_code.co_name - - is_handler = False - if isinstance(last_result, (BreakSignal, ContinueSignal)): - if func_name in ("visitWhileStmt", "visitForStmt"): - is_handler = True - elif isinstance(last_result, ReturnSignal): - if func_name in ("visitCallExp", "visitStart"): - is_handler = True - - if not is_handler: - stack.pop().close() - continue - - res = stack[-1].send(last_result) - - if hasattr(res, 'accept'): - next_gen = res.accept(self) - if inspect.isgenerator(next_gen): - stack.append(next_gen) - last_result = None - else: - last_result = next_gen + if signal is not None: + pending, signal = signal, None + res = stack[-1].throw(pending) else: - last_result = res - + res = stack[-1].send(last_result) except StopIteration as e: stack.pop() - last_result = e.value + res = e.value + except ControlSignal as e: + stack.pop() + signal = e + continue + else: + if hasattr(res, "accept"): + res = res.accept(self) + if inspect.isgenerator(res): + stack.append(res) + last_result = None + continue - if isinstance(last_result, ReturnSignal): - return last_result.value + if isinstance(res, ControlSignal): + signal = res + else: + last_result = res + + if isinstance(signal, ReturnSignal): + return signal.value + if signal is not None: + raise ValueError("break/continue used outside of a loop") return last_result def _is_tensor(self, val): return isinstance(val, torch.Tensor) or getattr(val, "is_nested", False) @@ -225,6 +361,16 @@ class UnifiedMathVisitor(MathExprVisitor): return ChangeTrackingDict.from_result(result_dict, tracking_source) return MathDict(result_dict) + def _truthy(self, value): + """Truth of a branch condition: a tensor or list is true when any element is non-zero.""" + if self._is_nested_tensor(value): + return any(self._truthy(t) for t in value.tensors) + if self._is_plain_tensor(value): + return bool(torch.any(value != 0)) + if self._is_list(value): + return any(self._truthy(v) for v in value) + return bool(value) + def _where_value(self, condition, true_value, false_value): if self._is_tensor(condition) or self._is_tensor(true_value) or self._is_tensor(false_value): condition_tensor = self._promote_to_tensor(condition) @@ -857,22 +1003,14 @@ class UnifiedMathVisitor(MathExprVisitor): if self._is_tensor(condition): true_val = yield ctx.expr(0) false_val = yield ctx.expr(1) - true_t = self._promote_to_tensor(true_val) - false_t = self._promote_to_tensor(false_val) - - if true_t.dtype != false_t.dtype: - true_t = true_t.float() - false_t = false_t.float() - - cond_t = torch.isclose(condition.float(), torch.tensor(0.0, device=self.device)) == False - return torch.where(cond_t, true_t, false_t).contiguous() + return self._where_value(condition, true_val, false_val) if self._is_list(condition): res = [] cache_true = None cache_false = None - for i, c in enumerate(condition): - if c: + for c in condition: + if self._truthy(c): if cache_true is None: cache_true = yield ctx.expr(0) res.append(cache_true) @@ -882,7 +1020,7 @@ class UnifiedMathVisitor(MathExprVisitor): res.append(cache_false) return res - if condition: + if self._truthy(condition): return (yield ctx.expr(0)) else: return (yield ctx.expr(1)) @@ -1418,10 +1556,7 @@ class UnifiedMathVisitor(MathExprVisitor): if self._is_list(val): k_val = max(0, min(k_val, len(val))) - try: - return sorted(val, reverse=True)[:k_val] - except: - return val[:k_val] + return sorted(val, reverse=True)[:k_val] return val @@ -1456,10 +1591,7 @@ class UnifiedMathVisitor(MathExprVisitor): if self._is_list(val): k_val = max(0, min(k_val, len(val))) - try: - return sorted(val)[:k_val] - except: - return val[:k_val] + return sorted(val)[:k_val] return val @@ -2899,14 +3031,17 @@ class UnifiedMathVisitor(MathExprVisitor): count = ctx.getChildCount() last_res = None - for i in range(count): - child = ctx.getChild(i) - if isinstance(child, TerminalNode): - continue + try: + for i in range(count): + child = ctx.getChild(i) + if isinstance(child, TerminalNode): + continue - res = yield child - if res is not None: - last_res = res + res = yield child + if res is not None: + last_res = res + except ReturnSignal as signal: + return signal.value return last_res @@ -2939,14 +3074,7 @@ class UnifiedMathVisitor(MathExprVisitor): def visitIfStmt(self, ctx): cond = yield ctx.expr() - def truthy(x): - if isinstance(x, torch.Tensor): - return torch.any(x != 0).item() - if isinstance(x, (list, tuple)): - return any(x) - return bool(x) - - if truthy(cond): + if self._truthy(cond): return (yield ctx.stmt(0)) elif ctx.stmt(1): return (yield ctx.stmt(1)) @@ -2955,24 +3083,14 @@ class UnifiedMathVisitor(MathExprVisitor): def visitWhileStmt(self, ctx): while True: throw_exception_if_processing_interrupted() - cond = yield ctx.expr() - is_true = ( - torch.any(cond != 0).item() if isinstance(cond, torch.Tensor) - else any(cond) if isinstance(cond, (list, tuple)) - else bool(cond) - ) - - if not is_true: + if not self._truthy((yield ctx.expr())): break - - res = yield ctx.stmt() - - if isinstance(res, ReturnSignal): - return res - if isinstance(res, BreakSignal): + try: + yield ctx.stmt() + except BreakSignal: break - if isinstance(res, ContinueSignal): - continue + except ContinueSignal: + pass return None def visitForStmt(self, ctx): @@ -2993,12 +3111,12 @@ class UnifiedMathVisitor(MathExprVisitor): for val in iterator: throw_exception_if_processing_interrupted() self.variables[var_name] = val - res = yield ctx.stmt() - - if isinstance(res, BreakSignal): + try: + yield ctx.stmt() + except BreakSignal: break - if isinstance(res, ContinueSignal): - continue + except ContinueSignal: + pass return None def visitBreakStmt(self, ctx): @@ -3108,8 +3226,8 @@ class UnifiedMathVisitor(MathExprVisitor): target[idx_tuple] = val_t return assigned_val if assign_op == "=" else val_t - except Exception as e: - raise ValueError(f"{ctx.start.line}:{ctx.start.column}: Indexed assignment to '{var_name}' failed: {str(e)}") + except (RuntimeError, IndexError, TypeError, ValueError) as e: + raise ValueError(f"{ctx.start.line}:{ctx.start.column}: Indexed assignment to '{var_name}' failed: {e}") from e elif self._is_list(target): # Recurse through nested lists if multiple indices provided curr = target @@ -3213,10 +3331,11 @@ class UnifiedMathVisitor(MathExprVisitor): self.depth += 1 self.variables["depth"] = float(self.depth) try: - res = yield lam.body - if isinstance(res, ReturnSignal): - return res.value - return res + return (yield lam.body) + except ReturnSignal as signal: + return signal.value + except (BreakSignal, ContinueSignal): + raise ValueError(f"{ctx.start.line}:{ctx.start.column}: break/continue outside of a loop in lambda body") finally: self.variables = self._scope_stack.pop() self.depth -= 1 @@ -3250,13 +3369,11 @@ class UnifiedMathVisitor(MathExprVisitor): self.variables["depth"] = float(self.depth) try: - # Visit the body using yield (trampoline will handle it) - res = yield func_def["body"] - if isinstance(res, ReturnSignal): - return res.value - - return res - + return (yield func_def["body"]) + except ReturnSignal as signal: + return signal.value + except (BreakSignal, ContinueSignal): + raise ValueError(f"{ctx.start.line}:{ctx.start.column}: break/continue outside of a loop in function '{func_name}'") finally: # Restore variables and depth self.variables = self._scope_stack.pop() @@ -3283,15 +3400,6 @@ class UnifiedMathVisitor(MathExprVisitor): return float(v.item()) return float(v) - def _to_bool(v): - if isinstance(v, bool): - return v - if self._is_tensor(v): - v = v.item() - if isinstance(v, str): - return v.lower() in ("true", "1", "yes") - return bool(v) - text = _to_str((yield ctx.expr(0))) font_name = _to_str((yield ctx.expr(1))) size = max(1, self._to_int((yield ctx.expr(2)), ctx,"text_image", strict=True)) @@ -3299,378 +3407,65 @@ class UnifiedMathVisitor(MathExprVisitor): weight_val = self._to_int((yield ctx.expr(4)), ctx,"text_image", strict=True) if len(ctx.expr()) > 4 else 400 angle = _to_float((yield ctx.expr(5))) if len(ctx.expr()) > 5 else 0.0 spacing = max(0.1, _to_float((yield ctx.expr(6)))) if len(ctx.expr()) > 6 else 1.0 - is_italic = _to_bool((yield ctx.expr(7))) if len(ctx.expr()) > 7 else False - has_ul = _to_bool((yield ctx.expr(8))) if len(ctx.expr()) > 8 else False + is_italic = self._to_bool((yield ctx.expr(7)), ctx, "text_image italic") if len(ctx.expr()) > 7 else False + has_ul = self._to_bool((yield ctx.expr(8)), ctx, "text_image underline") if len(ctx.expr()) > 8 else False - # --- Plně granulární mapování tlouštěk (CSS standard) --- - if weight_val <= 150: - weight_tier = "thin" - elif weight_val <= 250: - weight_tier = "extralight" - elif weight_val <= 350: - weight_tier = "light" - elif weight_val <= 450: - weight_tier = "regular" - elif weight_val <= 550: - weight_tier = "medium" - elif weight_val <= 650: - weight_tier = "semibold" - elif weight_val <= 750: - weight_tier = "bold" - elif weight_val <= 850: - weight_tier = "extrabold" + tier = _FONT_WEIGHT_TIERS[bisect.bisect_left((150, 250, 350, 450, 550, 650, 750, 850), weight_val)] + synthetic_italic = synthetic_bold = False + variation_weight, variation_italic = weight_val, is_italic + if os.path.isfile(font_name): + font_path = font_name else: - weight_tier = "black" - - if is_italic: - style_key = f"{weight_tier}_italic" if weight_tier != "regular" else "italic" - else: - style_key = weight_tier - - # --- Load font --- - def _normalise_font_name(value): - return value.lower().replace(" ", "").replace("-", "").replace("_", "") - - def _axis_name(axis): - name = axis.get("name", "") - if isinstance(name, bytes): - name = name.decode(errors="ignore") - return str(name).lower() - - def _axis_value(axis, value): - minimum = axis.get("minimum", value) - maximum = axis.get("maximum", value) - return max(minimum, min(maximum, value)) - - def _apply_font_variations(fnt, exact_weight, italic): - if not ( - hasattr(fnt, "get_variation_axes") - and hasattr(fnt, "set_variation_by_axes") - ): - return fnt - - try: - axes = fnt.get_variation_axes() - values = [] - changed = False - - for axis in axes: - axis_name = _axis_name(axis) - default = axis.get("default", axis.get("minimum", 0)) - - if axis_name in ("weight", "wght"): - values.append(_axis_value(axis, exact_weight)) - changed = True - elif axis_name in ("italic", "ital"): - values.append(_axis_value(axis, 1 if italic else 0)) - changed = True - elif axis_name in ("slant", "slnt"): - target = axis.get("minimum", -12) if italic else 0 - values.append(_axis_value(axis, target)) - changed = True - else: - values.append(default) - - if changed: - fnt.set_variation_by_axes(values) - except Exception: - pass - - return fnt - - def _load_font(name, sz, target_style, exact_weight): - if os.path.exists(name): - try: - fnt = ImageFont.truetype(name, sz) - return _apply_font_variations( - fnt, - exact_weight, - "italic" in target_style, - ) - except Exception: - pass - - style_suffixes = { - "thin": ["thin", "100", "hairline"], - "extralight": ["extralight", "ultralight", "200"], - "light": ["light", "300"], - "regular": ["regular", "normal", "reg", "400", "standard"], - "medium": ["medium", "500"], - "semibold": ["semibold", "demibold", "600"], - "bold": ["bold", "bd", "700"], - "extrabold": ["extrabold", "ultrabold", "800"], - "black": ["black", "heavy", "blk", "900"], - "italic": ["italic", "it", "i"], - }.get(target_style.replace("_italic", ""), [""]) - - if "italic" in target_style: - italic_suffixes = [] - for suffix in style_suffixes: - italic_suffixes.extend([ - f"{suffix}italic", - f"{suffix}ital", - f"{suffix}it", - f"{suffix}oblique", - ]) - style_suffixes = italic_suffixes + [ - "italic", "ital", "oblique", "it" - ] - - search_dirs = [] - import platform - sys_platform = platform.system() - if sys_platform == "Windows": - search_dirs = [ - os.path.join(os.environ.get("WINDIR", "C:\\Windows"), "Fonts"), - os.path.expandvars(r"%LOCALAPPDATA%\Microsoft\Windows\Fonts") - ] - elif sys_platform == "Darwin": - search_dirs = ["/Library/Fonts", "/System/Library/Fonts", os.path.expanduser("~/Library/Fonts")] - else: - xdg_data_home = os.environ.get("XDG_DATA_HOME", "~/.local/share") - search_dirs = [ - "/usr/share/fonts", - "/usr/local/share/fonts", - os.path.join(os.path.expanduser(xdg_data_home), "fonts"), - os.path.expanduser("~/.local/share/fonts"), - os.path.expanduser("~/.fonts"), - ] - - name_clean = _normalise_font_name(name) - wants_italic = "italic" in target_style - candidates = [] - - def _font_traits(base_name, sub_name): - combined = f"{base_name} {sub_name}" - is_italic_font = any( - token in combined - for token in ("italic", "ital", "oblique") - ) - is_italic_font = is_italic_font or base_name.endswith( - (f"{name_clean}i", f"{name_clean}bi") - ) - - weight = 400 - weight_hints = ( - (900, ("black", "heavy", "blk", "900")), - (800, ("extrabold", "ultrabold", "800")), - (700, ("bold", "bd", "700")), - (600, ("semibold", "demibold", "600")), - (500, ("medium", "500")), - (300, ("light", "300")), - (200, ("extralight", "ultralight", "200")), - (100, ("thin", "hairline", "100")), - ) - for font_weight, hints in weight_hints: - if any(hint in combined for hint in hints): - weight = font_weight - break - if base_name.endswith((f"{name_clean}bd", f"{name_clean}bi")): - weight = 700 - - return is_italic_font, weight - - def _candidate_score(base_name, sub_name, has_weight_axis): - is_italic_font, font_weight = _font_traits(base_name, sub_name) - italic_penalty = 0 if is_italic_font == wants_italic else 10000 - weight_penalty = 0 if has_weight_axis else abs(font_weight - exact_weight) - exact_name_penalty = 0 if base_name == name_clean else 50 - return italic_penalty + weight_penalty + exact_name_penalty - - for directory in search_dirs: - if not os.path.isdir(directory): - continue - for root, _, files in os.walk(directory): - for fname in files: - if not fname.lower().endswith((".ttf", ".otf", ".ttc", ".otc")): - continue - - full_path = os.path.join(root, fname) - base_name = _normalise_font_name(os.path.splitext(fname)[0]) - - match_found = False - has_weight_axis = False - sub_clean = "" - - # 1. KONTROLA NÁZVU SOUBORU - if name_clean in base_name: - if target_style == "regular" and base_name == name_clean: - match_found = True - else: - match_found = any(suf in base_name for suf in style_suffixes) - # 2. KONTROLA METADAT UVNITŘ SOUBORU - if not match_found: - try: - test_fnt = ImageFont.truetype(full_path, 10) - internal_family, internal_sub = test_fnt.getname() - - family_clean = _normalise_font_name(internal_family) - sub_clean = _normalise_font_name(internal_sub) - - if name_clean in family_clean or family_clean in name_clean: - has_variable_axes = False - if hasattr(test_fnt, "get_variation_axes"): - try: - axes = test_fnt.get_variation_axes() - axis_names = { - _axis_name(axis) - for axis in axes - } - has_weight_axis = bool( - axis_names & {"weight", "wght"} - ) - has_italic_axis = bool( - axis_names & { - "italic", - "ital", - "slant", - "slnt", - } - ) - has_variable_axes = ( - has_weight_axis - or has_italic_axis - ) - - if wants_italic and has_italic_axis: - match_found = True - elif not wants_italic and has_weight_axis: - match_found = True - except Exception: - has_variable_axes = False - - if not match_found and not has_variable_axes: - if target_style == "regular": - match_found = any( - s in sub_clean - for s in [ - "regular", - "normal", - "standard", - ] - ) or sub_clean == "" - else: - match_found = any( - suf in sub_clean - for suf in style_suffixes - ) or any( - suf in family_clean - for suf in style_suffixes - ) - except Exception: - continue - - if match_found: - candidates.append(( - _candidate_score( - base_name, - sub_clean, - has_weight_axis, - ), - full_path, - )) - - for _, full_path in sorted(candidates, key=lambda item: item[0]): - try: - fnt = ImageFont.truetype(full_path, sz) - return _apply_font_variations( - fnt, - exact_weight, - wants_italic, - ) - except Exception: - pass - - raise FileNotFoundError( - f"Font '{name}' with style '{target_style}' (weight {exact_weight}) was not found in the system directories." - ) - - try: - font = _load_font(font_name, size, style_key, weight_val) - synthetic_italic = False - synthetic_bold = False - except FileNotFoundError: - try: - font = _load_font(font_name, size, "regular", 400) + font_path = _find_font(font_name, tier, weight_val, is_italic) + if font_path is None: + # Fall back to the regular face and fake the requested style + font_path = _find_font(font_name, "regular", 400, False) + if font_path is None: + raise FileNotFoundError(f"Font '{font_name}' was not found in the system font directories.") synthetic_italic = is_italic - synthetic_bold = (weight_val >= 600) - except FileNotFoundError: - raise FileNotFoundError(f"Font '{font_name}' was not found in the system directories (nor its base version).") + synthetic_bold = weight_val >= 600 + variation_weight, variation_italic = 400, False + font = _apply_font_variations(ImageFont.truetype(font_path, size), variation_weight, variation_italic) # --- Word-wrap if max_width is set --- - def _wrap_text(text_in, fnt, max_w): - if max_w <= 0: - return text_in.splitlines() - lines = [] - for paragraph in text_in.splitlines(): - words = paragraph.split(" ") - current = "" - for word in words: - test = (current + " " + word).strip() - bbox = fnt.getbbox(test) - w = bbox[2] - bbox[0] - if w <= max_w or not current: - current = test - else: - lines.append(current) - current = word - lines.append(current) - return lines + lines = [] + for paragraph in text.splitlines(): + if max_width <= 0: + lines.append(paragraph) + continue + current = "" + for word in paragraph.split(" "): + test = (current + " " + word).strip() + if not current or font.getlength(test) <= max_width: + current = test + else: + lines.append(current) + current = word + lines.append(current) - lines = _wrap_text(text, font, max_width) + # --- Measure and render --- + stroke = max(1, int(size * 0.02)) if synthetic_bold else 0 + shear = 0.25 if synthetic_italic else 0.0 + bboxes = [font.getbbox(line or " ", stroke_width=stroke) for line in lines] + line_stride = max(1, int(max((b[3] - b[1] for b in bboxes), default=size) * spacing)) + canvas_h = max(1, line_stride * len(lines)) + canvas_w = max(1, max((b[2] - b[0] for b in bboxes), default=size) + int(shear * canvas_h)) - # --- Measure canvas size --- - line_heights = [] - line_widths = [] - for line in lines: - bbox = font.getbbox(line if line else " ") - line_widths.append(bbox[2] - bbox[0]) - line_heights.append(bbox[3] - bbox[1]) - - line_h = max(line_heights) if line_heights else size - line_stride = max(1, int(line_h * spacing)) - - padding_x = int(size * 0.3) if synthetic_italic else 0 - canvas_w = (max(line_widths) if line_widths else size) + padding_x - canvas_h = line_stride * len(lines) - - canvas_w = max(canvas_w, 1) - canvas_h = max(canvas_h, 1) - - # --- Render text --- img = Image.new("L", (canvas_w, canvas_h), color=0) draw = ImageDraw.Draw(img) - y = 0 - ul_thickness = max(1, int(size * 0.06)) ul_offset = max(1, int(size * 0.08)) - - for line in lines: + for i, (line, bbox) in enumerate(zip(lines, bboxes)): if not line: - y += line_stride continue - bbox = font.getbbox(line) - text_w = bbox[2] - bbox[0] - text_h = bbox[3] - bbox[1] - - # Synthetic bold. - if synthetic_bold: - thickness_offset = max(1, int(size * 0.02)) - for ox in range(-thickness_offset, thickness_offset + 1): - for oy in range(-thickness_offset, thickness_offset + 1): - draw.text((ox, (y - bbox[1]) + oy), line, fill=255, font=font) - else: - draw.text((0, y - bbox[1]), line, fill=255, font=font) - + y = i * line_stride + draw.text((-bbox[0], y - bbox[1]), line, fill=255, font=font, stroke_width=stroke, stroke_fill=255) if has_ul: - ul_top = y + text_h + ul_offset - draw.rectangle([0, ul_top, text_w, ul_top + ul_thickness], fill=255) + ul_top = y + bbox[3] - bbox[1] + ul_offset + draw.rectangle([0, ul_top, bbox[2] - bbox[0], ul_top + ul_thickness], fill=255) - y += line_stride - - # --- Software Italic (After Text Rendering) --- if synthetic_italic: - img = img.transform(img.size, Image.Transform.AFFINE, (1, 0.25, -0.25*text_h, 0, 1, 0)) + img = img.transform(img.size, Image.Transform.AFFINE, (1, shear, -shear * canvas_h, 0, 1, 0)) # --- Rotate if requested --- if angle != 0.0: @@ -3928,14 +3723,13 @@ class UnifiedMathVisitor(MathExprVisitor): def visitAnyFunc(self, ctx): val = yield ctx.expr() - if self._is_tensor(val): return torch.any(torch.isclose(val, torch.tensor(0.0, device=self.device)) == False).float() - if self._is_list(val): return float(any(val)) - return float(bool(val)) + if self._is_plain_tensor(val): return torch.any(val != 0).float() + return float(self._truthy(val)) def visitAllFunc(self, ctx): val = yield ctx.expr() - if self._is_tensor(val): return torch.all(torch.isclose(val, torch.tensor(0.0, device=self.device)) == False).float() - if self._is_list(val): return float(all(val)) + if self._is_plain_tensor(val): return torch.all(val != 0).float() + if self._is_list(val): return float(all(self._truthy(v) for v in val)) return float(bool(val)) def visitMedianFunc(self, ctx): @@ -4568,13 +4362,10 @@ class UnifiedMathVisitor(MathExprVisitor): def _bitwise_op(self, a, b, torch_op, scalar_op,ctx): """Binary bitwise operation handler supporting tensors, lists, and scalars.""" - try: - if self._is_tensor(a) and a.numel() == 1: - a = int(a.flatten()[0].item()) - if self._is_tensor(b) and b.numel() == 1: - b = int(b.flatten()[0].item()) - except Exception as e: - raise ValueError(f"Invalid tensor value for bitwise operation: {e}") + if self._is_tensor(a) and a.numel() == 1: + a = int(a.flatten()[0].item()) + if self._is_tensor(b) and b.numel() == 1: + b = int(b.flatten()[0].item()) # Handle tensor-list combinations if self._is_tensor(a) and self._is_list(b): @@ -5015,14 +4806,7 @@ class UnifiedMathVisitor(MathExprVisitor): if mask_idx < 0: return mask.view(new_shape).expand(overlay_shape) - # Fallback: try broadcast_to or expand - try: - return torch.broadcast_to(mask, overlay_shape) - except Exception: - try: - return mask.expand(overlay_shape) - except Exception: - return mask + return torch.broadcast_to(mask, overlay_shape) def visitOverlayFunc(self, ctx): base = yield ctx.expr(0) diff --git a/more_math/Parser/antlr_router.py b/more_math/Parser/antlr_router.py index 8fcada0..e865583 100644 --- a/more_math/Parser/antlr_router.py +++ b/more_math/Parser/antlr_router.py @@ -1,20 +1,8 @@ -import sys -import antlr4 +import importlib.metadata + def get_antlr_modules(): - runtime_version = "0.0.0" - - try: - if sys.version_info >= (3, 8): - import importlib.metadata - runtime_version = importlib.metadata.version("antlr4-python3-runtime") - else: - import pkg_resources - runtime_version = pkg_resources.get_distribution("antlr4-python3-runtime").version - except Exception: - runtime_version = getattr(antlr4, "__version__", "4.13.2") - - if runtime_version.startswith("4.9"): + if importlib.metadata.version("antlr4-python3-runtime").startswith("4.9"): from .legacy.MathExprLexer import MathExprLexer from .legacy.MathExprParser import MathExprParser from .legacy.MathExprVisitor import MathExprVisitor diff --git a/more_math/Parser/optical_flow_utils.py b/more_math/Parser/optical_flow_utils.py index 6fd4077..ce14ee4 100644 --- a/more_math/Parser/optical_flow_utils.py +++ b/more_math/Parser/optical_flow_utils.py @@ -4,134 +4,26 @@ import torch.nn.functional as F from torchvision.models.optical_flow import raft_large, Raft_Large_Weights from torchvision.utils import flow_to_image as torchvision_flow_to_image import numpy as np +import folder_paths +import comfy.utils _raft_model = None def get_raft_model(device): global _raft_model if _raft_model is None: - # First, try to load a user-provided checkpoint from ComfyUI's configurable - # "optical_flow" folder (preferred). If that fails, fall back to the - # torchvision default weights. - try: - # Imported here to avoid adding a hard dependency at module import time - # when this utility may be used outside of a full ComfyUI environment. - import folder_paths - import comfy.utils - - names = folder_paths.get_filename_list("optical_flow") - model_name = None - if names: - # Choose the first available file by default (matching other UI lists). - model_name = names[0] - else: - # Try to use torchvision weights but save them to the optical_flow folder - try: - weights = Raft_Large_Weights.DEFAULT - - # Instantiate the torchvision model (this will download weights into the torchvision cache) - try: - model_dl = raft_large(weights=weights, progress=True) - except Exception: - model_dl = raft_large(weights=weights, progress=False) - - # Move state dict to CPU and detach - state_dict = {k: v.cpu() for k, v in model_dl.state_dict().items()} - - # Try to infer original filename/extension from the weights metadata or URL - candidate_name = None - try: - url = getattr(weights, 'url', None) or getattr(weights, '_url', None) - if url: - candidate_name = os.path.basename(str(url).split('?')[0]) - except Exception: - candidate_name = None - if not candidate_name: - try: - meta = getattr(weights, 'meta', None) or {} - if isinstance(meta, dict): - candidate_name = meta.get('filename') or meta.get('name') - except Exception: - candidate_name = None - - # Ensure folder exists: use folder_paths mapping when available - save_dir = None - if hasattr(folder_paths, 'folder_names_and_paths') and 'optical_flow' in folder_paths.folder_names_and_paths: - dirs = folder_paths.folder_names_and_paths['optical_flow'][0] - if isinstance(dirs, (list, tuple)) and len(dirs) > 0: - save_dir = dirs[0] - if not save_dir: - save_dir = os.path.join(os.path.dirname(__file__), "..", "..", "models", "optical_flow") - os.makedirs(save_dir, exist_ok=True) - - # Decide filename and format - if candidate_name: - model_name = candidate_name - else: - model_name = "raft_large_default.pth" - - model_path = folder_paths.get_full_path("optical_flow", model_name) or os.path.join(save_dir, model_name) - - # Save in original format if known (.safetensors or .pth/.pt), otherwise default to .pth - _, ext = os.path.splitext(model_name) - ext = ext.lower() - saved_ok = False - try: - if ext in ('.safetensors', '.sft'): - # Save as safetensors - try: - comfy.utils.save_torch_file(state_dict, model_path) - saved_ok = True - except Exception: - import safetensors.torch - safetensors.torch.save_file(state_dict, model_path) - saved_ok = True - else: - # Default: save as a PyTorch checkpoint (.pth) - try: - import torch as _torch - _torch.save(state_dict, model_path) - saved_ok = True - except Exception: - # Fallback to comfy.utils (may write safetensors) - try: - comfy.utils.save_torch_file(state_dict, model_path) - saved_ok = True - except Exception: - saved_ok = False - except Exception: - saved_ok = False - - if not saved_ok: - model_name = None - except Exception: - model_name = None - - if model_name: - try: - model_path = folder_paths.get_full_path_or_raise("optical_flow", model_name) - sd = comfy.utils.load_torch_file(model_path, safe_load=True) - - # Basic validation for a RAFT-large state dict - has_raft_keys = ( - any(k.startswith("feature_encoder.") for k in sd) - and any(k.startswith("context_encoder.") for k in sd) - and any(k.startswith("update_block.") for k in sd) - ) - if has_raft_keys: - model = raft_large(weights=None, progress=False) - model.load_state_dict(sd) - _raft_model = model.to(device).eval() - return _raft_model - except Exception: - pass - except Exception: - # If anything goes wrong with folder resolution / loading, fall back below. - pass - - # Fallback to torchvision-provided default weights - weights = Raft_Large_Weights.DEFAULT - _raft_model = raft_large(weights=weights, progress=False).to(device).eval() + names = folder_paths.get_filename_list("optical_flow") + if names: + model = raft_large(weights=None, progress=False) + model.load_state_dict(comfy.utils.load_torch_file(folder_paths.get_full_path_or_raise("optical_flow", names[0]), safe_load=True)) + else: + # No local checkpoint: download torchvision's RAFT weights once and keep a copy in models/optical_flow + weights = Raft_Large_Weights.DEFAULT + model = raft_large(weights=weights, progress=True) + save_dir = folder_paths.get_folder_paths("optical_flow")[0] + os.makedirs(save_dir, exist_ok=True) + torch.save(model.state_dict(), os.path.join(save_dir, os.path.basename(weights.url))) + _raft_model = model.to(device).eval() return _raft_model def warp(x, flow, padding_mode='reflection'): diff --git a/more_math/SelectiveGuiderMathNode.py b/more_math/SelectiveGuiderMathNode.py index a5f6340..52b7f15 100644 --- a/more_math/SelectiveGuiderMathNode.py +++ b/more_math/SelectiveGuiderMathNode.py @@ -171,7 +171,7 @@ class SelectiveGuiderMathNode(io.ComfyNode): return None try: diffusion_model = patcher.get_model_object("diffusion_model") - except Exception: + except AttributeError: return None if stage in ("input", "output"): @@ -713,23 +713,16 @@ class SelectiveGuiderMathNode(io.ComfyNode): if hasattr(V, "model_patcher") and V.model_patcher is not None and not getattr(V, "_mrmth_debug_callbacks_added", False): def _on_register_all_hook_patches(model_patcher, hooks, target_dict, model_options, registered): - try: - reg_len = len(registered) if registered is not None else 0 - in_len = len(hooks) if hooks is not None else 0 - tr_count = len(registered.get_type(comfy.hooks.EnumHookType.TransformerOptions)) if registered is not None else 0 - dbg(f"REGISTER hooks_in={in_len} registered={reg_len} transformer_registered={tr_count}") - if model_options is not None: - rh = model_options.get("registered_hooks", None) - dbg(f"model_options.registered_hooks={'set' if rh is not None else 'None'}") - except Exception as e: - logging.warning(f"[SelectiveGuiderMathNode] REGISTER debug error: {e}") + reg_len = len(registered) if registered is not None else 0 + in_len = len(hooks) if hooks is not None else 0 + tr_count = len(registered.get_type(comfy.hooks.EnumHookType.TransformerOptions)) if registered is not None else 0 + dbg(f"REGISTER hooks_in={in_len} registered={reg_len} transformer_registered={tr_count}") + if model_options is not None: + rh = model_options.get("registered_hooks", None) + dbg(f"model_options.registered_hooks={'set' if rh is not None else 'None'}") def _on_apply_hooks(model_patcher, hooks): - try: - hl = len(hooks) if hooks is not None else 0 - dbg(f"APPLY hooks_len={hl}") - except Exception as e: - logging.warning(f"[SelectiveGuiderMathNode] APPLY debug error: {e}") + dbg(f"APPLY hooks_len={len(hooks) if hooks is not None else 0}") V.model_patcher.add_callback_with_key( comfy.patcher_extension.CallbacksMP.ON_REGISTER_ALL_HOOK_PATCHES, diff --git a/more_math/helper_functions.py b/more_math/helper_functions.py index ef77264..07ae26b 100644 --- a/more_math/helper_functions.py +++ b/more_math/helper_functions.py @@ -209,17 +209,12 @@ def move_to_device(value, device): return tuple(move_to_device(v, device) for v in value) if hasattr(value, "images") and hasattr(value, "audio"): - images = move_to_device(getattr(value, "images", None), device) - audio = move_to_device(getattr(value, "audio", None), device) - try: - return type(value)( - images=images, - audio=audio, - frame_rate=getattr(value, "frame_rate", None), - metadata=getattr(value, "metadata", None), - ) - except Exception: - return value + return type(value)( + images=move_to_device(value.images, device), + audio=move_to_device(value.audio, device), + frame_rate=getattr(value, "frame_rate", None), + metadata=getattr(value, "metadata", None), + ) return value @@ -320,22 +315,15 @@ def get_f_variable(f_dict): for k in sorted_keys: val = f_dict[k] if torch.is_tensor(val): - # If it's a tensor, just get the first element if it's a scalar or keep as is if it's already 1D? - # Usually F inputs are floats. If they are tensors, we'll try to treat them as values. - ordered_values.append(val.flatten()[0] if val.numel() > 0 else torch.tensor(0.0)) + ordered_values.append(float(val.flatten()[0]) if val.numel() > 0 else 0.0) elif isinstance(val, (int, float)): - ordered_values.append(torch.tensor(float(val))) + ordered_values.append(float(val)) else: - ordered_values.append(torch.tensor(0.0)) + ordered_values.append(0.0) if not ordered_values: - return None, 0 - - try: - stacked = torch.stack(ordered_values) - return stacked, len(ordered_values) - except Exception: - return None, len(ordered_values) + return None, 0 + return torch.tensor(ordered_values), len(ordered_values) def _collect_vars_from_node( node, needed_vars, assigned_vars, shadowed_vars): """Recursively collect variable reads from an AST node""" diff --git a/more_math/modelLikeCommon.py b/more_math/modelLikeCommon.py index 656c001..804a087 100644 --- a/more_math/modelLikeCommon.py +++ b/more_math/modelLikeCommon.py @@ -33,26 +33,6 @@ def get_effective_state_dict(obj): return None -def get_effective_weight(obj, key): - patcher = _get_model_patcher(obj) - if patcher is not None and hasattr(patcher, "patch_weight_to_device"): - try: - return patcher.patch_weight_to_device(key, return_weight=True) - except Exception: - pass - - if hasattr(obj, "patch_weight_to_device"): - try: - return obj.patch_weight_to_device(key, return_weight=True) - except Exception: - pass - - sd = get_effective_state_dict(obj) - if sd is not None: - return sd.get(key, None) - return None - - def get_base_weight(obj, key): """True pre-patch weight for `key`, regardless of whether the patcher currently has this key's patches baked into the live model (`backup` @@ -70,7 +50,7 @@ def get_base_weight(obj, key): try: weight, _, _ = comfy.model_patcher.get_key_weight(model, key) return weight - except Exception: + except AttributeError: pass sd = get_effective_state_dict(obj) @@ -90,10 +70,7 @@ def get_patched_weight(obj, key, base_weight): if patcher is not None: patches = getattr(patcher, "patches", None) if patches is not None and key in patches: - try: - return comfy.lora.calculate_weight(patches[key], base_weight.clone(), key) - except Exception: - pass + return comfy.lora.calculate_weight(patches[key], base_weight.clone(), key) return base_weight @@ -161,10 +138,6 @@ def calculate_patches_autogrow(Expr, V, F, pbar=None, mapping=None, stack=None, seen_keys.add(key) all_keys_list.append(key) - # Function to get weight from a valid object - def get_weight(obj, key): - return get_effective_weight(obj, key) - tree = None if isinstance(Expr,str): tree = parse_expr(Expr) diff --git a/tests/test_unified_visitor_regressions.py b/tests/test_unified_visitor_regressions.py index 3282959..401b138 100644 --- a/tests/test_unified_visitor_regressions.py +++ b/tests/test_unified_visitor_regressions.py @@ -80,7 +80,7 @@ def test_lambda_calls_require_exact_argument_count(call): def test_text_image_size_expression_is_evaluated_once(monkeypatch): visitor_module = importlib.import_module("more_math.Parser.UnifiedMathVisitor") default_font = ImageFont.load_default() - original_exists = os.path.exists + original_isfile = os.path.isfile monkeypatch.setattr( visitor_module.ImageFont, "truetype", @@ -88,8 +88,8 @@ def test_text_image_size_expression_is_evaluated_once(monkeypatch): ) monkeypatch.setattr( visitor_module.os.path, - "exists", - lambda path: path == "mock-font" or original_exists(path), + "isfile", + lambda path: path == "mock-font" or original_isfile(path), ) values = MathDict() @@ -249,3 +249,45 @@ def test_unconnected_stack_is_not_shared_between_runs(): _, stack = FloatMathNode.execute("stack_push(0, 1)", {"V0": 0.0}, remember_stack=True) assert stack == {0: [1.0]} + + +def test_top_level_return_stops_script(): + assert _visit("x = 1; return 5; x = 2; x", {}) == 5 + + +def test_return_from_loop_inside_function_resumes_caller(): + expression = "f(n) -> { for (i in [1, 2, 3]) { if (i == 2) return i * n; } return 0; }; f(10) + f(1)" + assert _visit(expression, {}) == 22 + + +def test_break_inside_function_does_not_escape_to_caller_loop(): + with pytest.raises(ValueError, match="outside of a loop"): + _visit("f() -> { break; }; for (i in [1, 2]) { f(); } 0", {}) + + +def test_tiny_non_zero_values_are_true(): + t = torch.tensor([1e-9, 0.0]) + + assert torch.equal(_visit("t ? 1 : 0", {"t": t}), torch.tensor([1.0, 0.0])) + assert _visit("any(t)", {"t": t}) == 1.0 + + +def test_text_image_synthetic_italic_with_only_empty_lines(monkeypatch): + visitor_module = importlib.import_module("more_math.Parser.UnifiedMathVisitor") + default_font = ImageFont.load_default() + monkeypatch.setattr(visitor_module.ImageFont, "truetype", lambda *args, **kwargs: default_font) + monkeypatch.setattr(visitor_module, "_find_font", lambda name, tier, weight, italic: None if italic else "mock-font") + + result = _visit('text_image("\\n", "MockFont", 12, 0, 400, 0, 1, 1)', {}) + + assert result.ndim == 2 + + +def test_bold_request_prefers_bold_over_semibold(monkeypatch): + visitor_module = importlib.import_module("more_math.Parser.UnifiedMathVisitor") + monkeypatch.setattr(visitor_module, "_system_fonts", lambda: [ + ("semibold.ttf", "arialsemibold", "arial", "semibold", False, False), + ("bold.ttf", "arialbold", "arial", "bold", False, False), + ]) + + assert visitor_module._find_font("Arial", "bold", 700, False) == "bold.ttf"