5359 lines
232 KiB
Python
5359 lines
232 KiB
Python
from __future__ import annotations
|
|
|
|
import gc
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
import re
|
|
import shlex
|
|
import subprocess
|
|
import threading
|
|
import time
|
|
from dataclasses import asdict, dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
CATEGORY = "prompt/diffusiongemma"
|
|
COMPATIBILITY_CATEGORY = f"{CATEGORY}/compatibility"
|
|
OPTIONAL_CATEGORY = f"{CATEGORY}/optional"
|
|
MODEL_TYPE = "DG_MODEL_CONFIG"
|
|
MEDIA_TYPE = "DG_MEDIA_CONTEXT"
|
|
CONTEXT_TYPE = "DG_CONTEXT"
|
|
TARGET_PROFILE_TYPE = "DG_TARGET_PROFILE_CONFIG"
|
|
|
|
DEFAULT_GGUF_MODEL_PATH = r"Z:\models\LLM\diffusiongemma-26B-A4B-it-GGUF\diffusiongemma-26B-A4B-it-Q4_K_M.gguf"
|
|
DEFAULT_GGUF_CLI_PATH = r"C:\tools\llama.cpp-diffusiongemma\build-clang-cuda\bin\llama-diffusion-cli.exe"
|
|
DEFAULT_GGUF_EXTRA_ARGS = "-ngl all -cmoe -sm none -mg 0 --diffusion-eb auto --diffusion-kv-cache auto --diffusion-gpu-sampling auto --diffusion-eb-max-steps 24 --diffusion-steps 24 --temp 0.45 -no-cnv --log-verbosity 1"
|
|
DEFAULT_CUDA_RUNTIME_BIN = r"C:\tools\cuda-13.3-conda\Library\bin"
|
|
DEFAULT_NVFP4_MODEL_DIR = str(Path.cwd() / "models" / "LLM" / "diffusiongemma-26B-A4B-it-NVFP4")
|
|
DEFAULT_TRANSFORMERS_MODEL_DIR = str(Path.cwd() / "models" / "LLM" / "diffusiongemma-26B-A4B-it")
|
|
DEFAULT_ADVANCED_MODEL_PATH = DEFAULT_NVFP4_MODEL_DIR if Path(DEFAULT_NVFP4_MODEL_DIR).exists() else DEFAULT_TRANSFORMERS_MODEL_DIR
|
|
DEFAULT_MODEL_PATH = DEFAULT_GGUF_MODEL_PATH if Path(DEFAULT_GGUF_MODEL_PATH).exists() else DEFAULT_TRANSFORMERS_MODEL_DIR
|
|
FINAL_JSON_OPEN = "<|final_json|>"
|
|
FINAL_JSON_CLOSE = "<|end_final_json|>"
|
|
FINAL_JSON_OPEN_MARKERS = (FINAL_JSON_OPEN, "<|answer|>", "<|final|>")
|
|
FINAL_JSON_CLOSE_MARKERS = (FINAL_JSON_CLOSE, "<|/final_json|>", "<|end_final|>", "<|end|>")
|
|
MIN_MODELOPT_NVFP4_BRIDGE_MEMORY_GB = 18.0
|
|
MEMORY_PRESET_SAFE_20 = "20.0 GB - safe practical"
|
|
MEMORY_PRESET_MIN_18 = "18.0 GB - minimum supported"
|
|
MEMORY_PRESET_FULL_GPU = "0 GB - full GPU"
|
|
MEMORY_PRESET_CHOICES = [MEMORY_PRESET_SAFE_20, MEMORY_PRESET_MIN_18, MEMORY_PRESET_FULL_GPU]
|
|
RESOLUTION_SELECTOR_ASPECT_LABELS = {
|
|
"1:1": "1:1 (Square)",
|
|
"2:3": "2:3 (Portrait Photo)",
|
|
"3:2": "3:2 (Photo)",
|
|
"3:4": "3:4 (Portrait Standard)",
|
|
"4:3": "4:3 (Standard)",
|
|
"9:16": "9:16 (Portrait Widescreen)",
|
|
"16:9": "16:9 (Widescreen)",
|
|
"21:9": "21:9 (Ultrawide)",
|
|
}
|
|
RESOLUTION_SELECTOR_RATIOS = {
|
|
"1:1": (1, 1),
|
|
"2:3": (2, 3),
|
|
"3:2": (3, 2),
|
|
"3:4": (3, 4),
|
|
"4:3": (4, 3),
|
|
"9:16": (9, 16),
|
|
"16:9": (16, 9),
|
|
"21:9": (21, 9),
|
|
}
|
|
REFERENCE_PREP_ASPECT_CHOICES = ["source", "1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3", "21:9"]
|
|
MEDIA_SYNTHESIS_MODE_CHOICES = ["video_recreation", "image_identity_video_control"]
|
|
DEFAULT_MASTER_PROMPT = """You are a local ComfyUI prompt architect. Treat the user's text as a content brief: it only needs to say what the video or image should be about. Do not require the user to ask for prompt-writing boilerplate such as detailed subject, background, lighting, camera movement, pacing, composition, or style; infer useful generator-ready detail when it is missing, while preserving every explicit user constraint.
|
|
|
|
Use the user's text, any backend-consumable pixels, and any supplied visual_description to produce final, practical prompts for generation models. If the active backend is metadata-only for image/video, do not pretend you saw pixels; use only the user's text, the visual_description, and factual media metadata. Preserve requested subjects, exact text, composition, motion, style, camera behavior, visual capture language, audio mode, audio guidance, and constraints. In image_identity_video_control workflows, treat a request to replace the video subject with the image subject as a strict identity-and-appearance transfer: face, hair, body type, wardrobe, accessories, styling, and distinguishing traits come from the reference image, while the video contributes pose, action, blocking, depth/canny/edge structure, camera choreography, timing, composition, and scene geometry. Do not copy the control-video subject's wardrobe, hair, face, body type, accessories, or styling into the final prompt unless the user explicitly asks for those video-subject traits. For LTX source-video work, explicitly observe camera choreography before writing the final prompt: first-frame and ending framing, shot size, camera height/angle, camera support or stability, camera path, zoom/dolly/truck/boom/pan/tilt/roll/orbit/rotation, subject-camera relationship, focus behavior, depth/parallax, and visible lens/capture changes. Transcribe observed camera work precisely instead of reducing it to generic phrases like camera movement, zoom out, or pan left. Visual capture language includes exposure level, overexposed or blown highlights, underexposure, flash/direct light, bloom, halation, lens blur, bokeh, depth of field, focal length impression, perspective compression, lens distortion, motion blur, grain/noise, white balance, contrast, dynamic range, and camera angle. Do not describe invisible intent. Every prompt field must contain the finished prompt text or finished prompt JSON, never instructions about how to write a prompt. Return only valid JSON with ltx_prompt, ideogram_prompt, negative_prompt, scene_segments, and metadata."""
|
|
|
|
|
|
@dataclass
|
|
class RuntimeConfig:
|
|
model_path: str
|
|
backend: str
|
|
dtype: str
|
|
quantization: str
|
|
local_files_only: bool
|
|
unload_policy: str
|
|
max_memory_gb: float
|
|
cli_path: str = ""
|
|
extra_args: str = ""
|
|
temperature: float = 0.45
|
|
fallback_backend: str = "template"
|
|
status: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class MediaContext:
|
|
images: Any = None
|
|
source: str = "none"
|
|
metadata: dict[str, Any] = field(default_factory=dict)
|
|
|
|
|
|
@dataclass
|
|
class GemmaContext:
|
|
user_prompt: str = ""
|
|
images: Any = None
|
|
source: str = "none"
|
|
media_metadata: dict[str, Any] = field(default_factory=dict)
|
|
visual_description: str = ""
|
|
warnings: list[str] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class TargetProfileConfig:
|
|
target_profile: str = "ltx"
|
|
audio_mode: str = "auto_scene_audio"
|
|
audio_guidance: str = ""
|
|
target_duration_seconds: float = 0.0
|
|
ltx_style: str = ""
|
|
ideogram_aspect_ratio: str = "1:1"
|
|
ideogram_render_style: str = ""
|
|
ideogram_exact_text: str = ""
|
|
ideogram_json_output: bool = True
|
|
negative_prompt_mode: str = "auto"
|
|
negative_prompt_guidance: str = ""
|
|
|
|
|
|
_MODEL_CACHE: dict[str, tuple[Any, Any]] = {}
|
|
_NVFP4_SMOKE_CACHE: dict[str, dict[str, Any]] = {}
|
|
|
|
|
|
def _dg_log(message: str, *args: Any) -> None:
|
|
logging.info("[DiffusionGemma] " + message, *args)
|
|
|
|
|
|
def _progress_enabled() -> bool:
|
|
return os.environ.get("DG_PROGRESS", "1").strip().lower() not in {"0", "false", "off", "no"}
|
|
|
|
|
|
def _release_transformers_runtime() -> None:
|
|
_MODEL_CACHE.clear()
|
|
gc.collect()
|
|
try:
|
|
import torch
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
try:
|
|
torch.cuda.ipc_collect()
|
|
except Exception:
|
|
pass
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def _json_dumps(data: Any) -> str:
|
|
return json.dumps(data, indent=2, sort_keys=True, ensure_ascii=True)
|
|
|
|
|
|
def _prompt_json_dumps(data: Any) -> str:
|
|
return json.dumps(data, indent=2, ensure_ascii=False)
|
|
|
|
|
|
def _ideogram_json_dumps(data: Any) -> str:
|
|
return json.dumps(data, separators=(",", ":"), ensure_ascii=False)
|
|
|
|
|
|
def _safe_float(value: Any, default: float = 0.0) -> float:
|
|
try:
|
|
return float(value)
|
|
except Exception:
|
|
return default
|
|
|
|
|
|
def _memory_preset_to_gb(value: Any) -> float:
|
|
if isinstance(value, (int, float)):
|
|
return float(value)
|
|
text = str(value or "").strip()
|
|
lowered = text.lower()
|
|
if "full" in lowered:
|
|
return 0.0
|
|
if "minimum" in lowered:
|
|
return MIN_MODELOPT_NVFP4_BRIDGE_MEMORY_GB
|
|
if "safe" in lowered:
|
|
return 20.0
|
|
match = re.search(r"[-+]?\d+(?:\.\d+)?", text)
|
|
if match:
|
|
return _safe_float(match.group(0), 20.0)
|
|
return 20.0
|
|
|
|
|
|
def _run_blocking_with_progress(
|
|
label: str,
|
|
fn: Any,
|
|
*,
|
|
node_id: str | None = None,
|
|
estimated_seconds: float | None = None,
|
|
) -> Any:
|
|
if not _progress_enabled():
|
|
return fn()
|
|
|
|
estimate = max(5.0, _safe_float(os.environ.get("DG_PROGRESS_ESTIMATE_SECONDS", ""), estimated_seconds or 120.0))
|
|
estimate_units = int(math.ceil(estimate))
|
|
interval = max(0.25, _safe_float(os.environ.get("DG_PROGRESS_INTERVAL_SECONDS", ""), 1.0))
|
|
result_box: dict[str, Any] = {}
|
|
|
|
def worker() -> None:
|
|
try:
|
|
result_box["value"] = fn()
|
|
except BaseException as exc:
|
|
result_box["error"] = exc
|
|
|
|
progress_bar = None
|
|
tqdm_bar = None
|
|
try:
|
|
from comfy.utils import ProgressBar
|
|
|
|
progress_bar = ProgressBar(estimate_units, node_id=node_id)
|
|
progress_bar.update_absolute(0, estimate_units)
|
|
except Exception:
|
|
progress_bar = None
|
|
try:
|
|
from tqdm.auto import tqdm
|
|
|
|
tqdm_bar = tqdm(
|
|
total=estimate_units,
|
|
desc=label,
|
|
unit="s",
|
|
leave=True,
|
|
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}]",
|
|
)
|
|
except Exception:
|
|
tqdm_bar = None
|
|
|
|
thread = threading.Thread(target=worker, name=f"{label} worker", daemon=True)
|
|
start = time.perf_counter()
|
|
last_reported = 0
|
|
displayed_total = estimate_units
|
|
thread.start()
|
|
try:
|
|
while thread.is_alive():
|
|
elapsed = time.perf_counter() - start
|
|
displayed_total = max(displayed_total, int(math.ceil(elapsed + interval)))
|
|
current = min(int(elapsed), displayed_total)
|
|
if progress_bar is not None:
|
|
progress_bar.update_absolute(current, displayed_total)
|
|
if tqdm_bar is not None:
|
|
if displayed_total != tqdm_bar.total:
|
|
tqdm_bar.total = displayed_total
|
|
tqdm_bar.refresh()
|
|
delta = max(0, current - last_reported)
|
|
if delta:
|
|
tqdm_bar.update(delta)
|
|
last_reported += delta
|
|
time.sleep(interval)
|
|
thread.join()
|
|
elapsed = time.perf_counter() - start
|
|
total = max(1, last_reported, int(math.ceil(elapsed)))
|
|
if progress_bar is not None:
|
|
progress_bar.update_absolute(total, total)
|
|
if tqdm_bar is not None:
|
|
if total != tqdm_bar.total:
|
|
tqdm_bar.total = total
|
|
remaining = max(0.0, total - tqdm_bar.n)
|
|
if remaining:
|
|
tqdm_bar.update(remaining)
|
|
finally:
|
|
if tqdm_bar is not None:
|
|
tqdm_bar.close()
|
|
|
|
if "error" in result_box:
|
|
raise result_box["error"]
|
|
return result_box.get("value")
|
|
|
|
|
|
def _normalize_aspect_ratio_key(aspect_ratio: Any) -> str:
|
|
text = str(aspect_ratio or "").strip()
|
|
match = re.search(r"\b(\d+)\s*:\s*(\d+)\b", text)
|
|
if not match:
|
|
return "1:1"
|
|
key = f"{int(match.group(1))}:{int(match.group(2))}"
|
|
return key if key in RESOLUTION_SELECTOR_RATIOS else "1:1"
|
|
|
|
|
|
def _resolution_selector_preset(aspect_ratio: Any) -> str:
|
|
return RESOLUTION_SELECTOR_ASPECT_LABELS[_normalize_aspect_ratio_key(aspect_ratio)]
|
|
|
|
|
|
def _known_aspect_ratio_label_from_dimensions(width: int, height: int) -> str:
|
|
src_w = max(1, int(width))
|
|
src_h = max(1, int(height))
|
|
source_ratio = src_w / float(src_h)
|
|
best_key = ""
|
|
best_delta = float("inf")
|
|
for key, (w_ratio, h_ratio) in RESOLUTION_SELECTOR_RATIOS.items():
|
|
delta = abs(source_ratio - (w_ratio / float(h_ratio)))
|
|
if delta < best_delta:
|
|
best_delta = delta
|
|
best_key = key
|
|
if best_key and best_delta <= 0.02:
|
|
return RESOLUTION_SELECTOR_ASPECT_LABELS[best_key]
|
|
return f"Custom / source ({src_w}x{src_h})"
|
|
|
|
|
|
def _resolution_selector_dimensions(aspect_ratio: Any, megapixels: Any = 2.0, multiple: Any = 8) -> tuple[int, int]:
|
|
key = _normalize_aspect_ratio_key(aspect_ratio)
|
|
w_ratio, h_ratio = RESOLUTION_SELECTOR_RATIOS[key]
|
|
mp = max(0.1, min(16.0, _safe_float(megapixels, 2.0)))
|
|
mult = int(max(8, min(128, _safe_float(multiple, 8))))
|
|
mult = max(1, int(round(mult / 4)) * 4)
|
|
total_pixels = mp * 1024 * 1024
|
|
scale = math.sqrt(total_pixels / (w_ratio * h_ratio))
|
|
width = max(mult, round(w_ratio * scale / mult) * mult)
|
|
height = max(mult, round(h_ratio * scale / mult) * mult)
|
|
return int(width), int(height)
|
|
|
|
|
|
def _round_dimension_to_multiple(value: float, multiple: int) -> int:
|
|
mult = int(max(1, multiple))
|
|
return max(mult, int(round(float(value) / mult) * mult))
|
|
|
|
|
|
def _dimensions_for_long_edge(width: int, height: int, long_edge: int, multiple: int = 8) -> tuple[int, int]:
|
|
src_w = max(1, int(width))
|
|
src_h = max(1, int(height))
|
|
edge = max(64, int(long_edge))
|
|
scale = edge / float(max(src_w, src_h))
|
|
return (
|
|
_round_dimension_to_multiple(src_w * scale, multiple),
|
|
_round_dimension_to_multiple(src_h * scale, multiple),
|
|
)
|
|
|
|
|
|
def _aspect_dimensions_for_long_edge(aspect_ratio: Any, long_edge: int, multiple: int = 8) -> tuple[int, int]:
|
|
key = _normalize_aspect_ratio_key(aspect_ratio)
|
|
w_ratio, h_ratio = RESOLUTION_SELECTOR_RATIOS[key]
|
|
edge = max(64, int(long_edge))
|
|
if w_ratio >= h_ratio:
|
|
width = edge
|
|
height = max(1, int(round(edge * (h_ratio / float(w_ratio)))))
|
|
else:
|
|
height = edge
|
|
width = max(1, int(round(edge * (w_ratio / float(h_ratio)))))
|
|
return (_round_dimension_to_multiple(width, multiple), _round_dimension_to_multiple(height, multiple))
|
|
|
|
|
|
def _center_crop_dimensions_for_aspect(width: int, height: int, aspect_ratio: Any) -> tuple[int, int]:
|
|
key = _normalize_aspect_ratio_key(aspect_ratio)
|
|
target_w, target_h = RESOLUTION_SELECTOR_RATIOS[key]
|
|
src_w = max(1, int(width))
|
|
src_h = max(1, int(height))
|
|
src_ratio = src_w / float(src_h)
|
|
dst_ratio = target_w / float(target_h)
|
|
if abs(src_ratio - dst_ratio) < 1e-6:
|
|
return (src_w, src_h)
|
|
if src_ratio > dst_ratio:
|
|
crop_w = max(1, int(round(src_h * dst_ratio)))
|
|
return (min(crop_w, src_w), src_h)
|
|
crop_h = max(1, int(round(src_w / dst_ratio)))
|
|
return (src_w, min(crop_h, src_h))
|
|
|
|
|
|
def _center_crop_image_batch(image: Any, crop_width: int, crop_height: int) -> Any:
|
|
try:
|
|
import torch
|
|
|
|
tensor = image if isinstance(image, torch.Tensor) else torch.as_tensor(image)
|
|
if tensor.ndim == 3:
|
|
tensor = tensor.unsqueeze(0)
|
|
_, src_h, src_w, _ = tensor.shape
|
|
crop_w = max(1, min(int(crop_width), int(src_w)))
|
|
crop_h = max(1, min(int(crop_height), int(src_h)))
|
|
x0 = max(0, (int(src_w) - crop_w) // 2)
|
|
y0 = max(0, (int(src_h) - crop_h) // 2)
|
|
return tensor[:, y0 : y0 + crop_h, x0 : x0 + crop_w, :]
|
|
except Exception:
|
|
return image
|
|
|
|
|
|
def _resize_image_batch(image: Any, width: int, height: int) -> Any:
|
|
import torch
|
|
|
|
tensor = image if isinstance(image, torch.Tensor) else torch.as_tensor(image)
|
|
if tensor.ndim == 3:
|
|
tensor = tensor.unsqueeze(0)
|
|
original_dtype = tensor.dtype
|
|
original_device = tensor.device
|
|
work = tensor
|
|
if not torch.is_floating_point(work):
|
|
work = work.float()
|
|
nchw = work.permute(0, 3, 1, 2).contiguous()
|
|
try:
|
|
resized = torch.nn.functional.interpolate(
|
|
nchw,
|
|
size=(max(1, int(height)), max(1, int(width))),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
antialias=True,
|
|
)
|
|
except TypeError:
|
|
resized = torch.nn.functional.interpolate(
|
|
nchw,
|
|
size=(max(1, int(height)), max(1, int(width))),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
result = resized.permute(0, 2, 3, 1).contiguous()
|
|
if torch.is_floating_point(tensor) and result.dtype != original_dtype:
|
|
result = result.to(dtype=original_dtype)
|
|
return result.to(device=original_device)
|
|
|
|
|
|
def _read_json_file(path: Path) -> dict[str, Any]:
|
|
try:
|
|
data = json.loads(path.read_text(encoding="utf-8"))
|
|
return data if isinstance(data, dict) else {}
|
|
except Exception:
|
|
return {}
|
|
|
|
|
|
def _pad_nvfp4_block_scale_rows(block_scale: Any, output_rows: int) -> Any:
|
|
import torch
|
|
|
|
rounded_rows = ((int(output_rows) + 127) // 128) * 128
|
|
if int(block_scale.shape[0]) >= rounded_rows:
|
|
return block_scale
|
|
pad_shape = (rounded_rows - int(block_scale.shape[0]), int(block_scale.shape[1]))
|
|
padding = torch.zeros(pad_shape, dtype=block_scale.dtype, device=block_scale.device)
|
|
return torch.cat([block_scale, padding], dim=0)
|
|
|
|
|
|
def _modelopt_to_comfy_fp4_packing(qdata: Any) -> Any:
|
|
return ((qdata & 0x0F) << 4) | ((qdata & 0xF0) >> 4)
|
|
|
|
|
|
def _modelopt_to_comfy_block_scale(block_scale: Any) -> Any:
|
|
from comfy_kitchen.float_utils import to_blocked
|
|
|
|
return to_blocked(block_scale, flatten=False)
|
|
|
|
|
|
def _safetensors_index(model_path: str) -> dict[str, str]:
|
|
index_path = Path(model_path) / "model.safetensors.index.json"
|
|
index_data = _read_json_file(index_path)
|
|
weight_map = index_data.get("weight_map")
|
|
return weight_map if isinstance(weight_map, dict) else {}
|
|
|
|
|
|
def _nvfp4_bridge_smoke(model_path: str) -> dict[str, Any]:
|
|
cache_key = str(Path(model_path))
|
|
cached = _NVFP4_SMOKE_CACHE.get(cache_key)
|
|
if cached is not None:
|
|
return cached.copy()
|
|
result: dict[str, Any] = {"passed": False, "error": "", "projection": ""}
|
|
try:
|
|
import torch
|
|
from safetensors import safe_open
|
|
from comfy.quant_ops import QuantizedTensor, get_layout_class
|
|
|
|
if not torch.cuda.is_available():
|
|
raise RuntimeError("CUDA is not available.")
|
|
root = Path(model_path)
|
|
weight_map = _safetensors_index(model_path)
|
|
prefix = "model.decoder.layers.0.experts.0.gate_proj"
|
|
required = [f"{prefix}.{name}" for name in ("weight", "weight_scale", "weight_scale_2", "input_scale")]
|
|
if not all(key in weight_map for key in required):
|
|
raise RuntimeError("Expected layer-0 expert gate_proj NVFP4 tensors were not found.")
|
|
|
|
handles: dict[str, Any] = {}
|
|
|
|
def get_tensor(key: str) -> Any:
|
|
shard = weight_map[key]
|
|
handle = handles.get(shard)
|
|
if handle is None:
|
|
handle = safe_open(str(root / shard), framework="pt", device="cpu")
|
|
handles[shard] = handle
|
|
return handle.get_tensor(key)
|
|
|
|
try:
|
|
qdata = _modelopt_to_comfy_fp4_packing(get_tensor(f"{prefix}.weight")).to("cuda", non_blocking=True)
|
|
block_scale = _modelopt_to_comfy_block_scale(get_tensor(f"{prefix}.weight_scale")).to("cuda", non_blocking=True)
|
|
tensor_scale = get_tensor(f"{prefix}.weight_scale_2").to("cuda", non_blocking=True)
|
|
input_scale = get_tensor(f"{prefix}.input_scale").to("cuda", non_blocking=True)
|
|
finally:
|
|
handles.clear()
|
|
|
|
out_features = int(qdata.shape[0])
|
|
in_features = int(qdata.shape[1]) * 2
|
|
layout_cls = get_layout_class("TensorCoreNVFP4Layout")
|
|
weight_params = layout_cls.Params(
|
|
scale=tensor_scale,
|
|
block_scale=block_scale,
|
|
orig_dtype=torch.bfloat16,
|
|
orig_shape=(out_features, in_features),
|
|
)
|
|
weight = QuantizedTensor(qdata.to(torch.uint8), "TensorCoreNVFP4Layout", weight_params)
|
|
x = torch.zeros((2, in_features), dtype=torch.bfloat16, device="cuda")
|
|
qx = QuantizedTensor.from_float(x, "TensorCoreNVFP4Layout", scale=input_scale)
|
|
y = torch.nn.functional.linear(qx, weight, None)
|
|
torch.cuda.synchronize()
|
|
if not bool(torch.isfinite(y.float()).all()):
|
|
raise RuntimeError("NVFP4 bridge smoke produced non-finite output.")
|
|
result.update(
|
|
{
|
|
"passed": True,
|
|
"projection": prefix,
|
|
"output_shape": list(y.shape),
|
|
"output_dtype": str(y.dtype).replace("torch.", ""),
|
|
}
|
|
)
|
|
except Exception as exc:
|
|
result["error"] = f"{type(exc).__name__}: {exc}"
|
|
finally:
|
|
gc.collect()
|
|
try:
|
|
import torch
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
except Exception:
|
|
pass
|
|
_NVFP4_SMOKE_CACHE[cache_key] = result.copy()
|
|
return result
|
|
|
|
|
|
def _model_path_info(model_path: str) -> dict[str, Any]:
|
|
path = Path(model_path or "")
|
|
info: dict[str, Any] = {
|
|
"path_kind": "missing",
|
|
"is_gguf_file": False,
|
|
"is_hf_repo": False,
|
|
"model_type": "",
|
|
"architecture": "",
|
|
"quant_algo": "",
|
|
"quant_method": "",
|
|
"quant_producer": "",
|
|
"quant_producer_version": "",
|
|
"transformers_version_required": "",
|
|
"has_generation_config": False,
|
|
"has_modelopt_state": False,
|
|
"runtime_hint": "",
|
|
}
|
|
if not path.exists():
|
|
return info
|
|
if path.is_file():
|
|
info["path_kind"] = "gguf_file" if path.suffix.lower() == ".gguf" else "file"
|
|
info["is_gguf_file"] = path.suffix.lower() == ".gguf"
|
|
if info["is_gguf_file"]:
|
|
info["runtime_hint"] = "llama-diffusion-cli GGUF"
|
|
return info
|
|
if not path.is_dir():
|
|
info["path_kind"] = "unknown"
|
|
return info
|
|
|
|
config = _read_json_file(path / "config.json")
|
|
hf_quant_config = _read_json_file(path / "hf_quant_config.json")
|
|
quant_config = config.get("quantization_config") if isinstance(config.get("quantization_config"), dict) else {}
|
|
if not quant_config:
|
|
quant_config = hf_quant_config.get("quantization") if isinstance(hf_quant_config.get("quantization"), dict) else {}
|
|
producer = quant_config.get("producer") if isinstance(quant_config.get("producer"), dict) else {}
|
|
if not producer:
|
|
producer = hf_quant_config.get("producer") if isinstance(hf_quant_config.get("producer"), dict) else {}
|
|
architectures = config.get("architectures") if isinstance(config.get("architectures"), list) else []
|
|
info.update(
|
|
{
|
|
"is_hf_repo": bool(config),
|
|
"model_type": str(config.get("model_type", "")),
|
|
"architecture": str(architectures[0]) if architectures else "",
|
|
"quant_algo": str(quant_config.get("quant_algo", "")),
|
|
"quant_method": str(quant_config.get("quant_method", config.get("quantization_config", {}).get("quant_method", "")) if isinstance(config.get("quantization_config"), dict) else ""),
|
|
"quant_producer": str(producer.get("name", "")),
|
|
"quant_producer_version": str(producer.get("version", "")),
|
|
"transformers_version_required": str(config.get("transformers_version", "")),
|
|
"has_generation_config": (path / "generation_config.json").exists(),
|
|
"has_modelopt_state": (path / "modelopt_state.pth").exists(),
|
|
}
|
|
)
|
|
if info["quant_algo"].upper() == "NVFP4":
|
|
info["path_kind"] = "nvfp4_hf_repo"
|
|
info["runtime_hint"] = "NVIDIA ModelOpt NVFP4 checkpoint"
|
|
elif config:
|
|
info["path_kind"] = "hf_repo"
|
|
info["runtime_hint"] = "Transformers whole-repo model"
|
|
else:
|
|
info["path_kind"] = "directory"
|
|
return info
|
|
|
|
|
|
def _runtime_status(config: RuntimeConfig) -> dict[str, Any]:
|
|
model_info = _model_path_info(config.model_path)
|
|
status: dict[str, Any] = {
|
|
"backend": config.backend,
|
|
"model_path": config.model_path,
|
|
"model_path_exists": Path(config.model_path).exists(),
|
|
**model_info,
|
|
"local_files_only": config.local_files_only,
|
|
"max_memory_gb": config.max_memory_gb,
|
|
"temperature": config.temperature,
|
|
"transformers_importable": False,
|
|
"diffusion_gemma_module": False,
|
|
"diffusion_gemma_class": False,
|
|
"auto_processor": False,
|
|
"vllm_importable": False,
|
|
"modelopt_importable": False,
|
|
"comfy_kitchen_importable": False,
|
|
"fp_quant_importable": False,
|
|
"qutlass_importable": False,
|
|
"tensorrt_llm_importable": False,
|
|
"torch_cuda_available": False,
|
|
"torch_cuda_capability": "",
|
|
"transformers_version": "",
|
|
"processor_class": "",
|
|
"modelopt_nvfp4_supported": False,
|
|
"comfy_nvfp4_bridge_supported": False,
|
|
"nvfp4_bridge_min_memory_gb": MIN_MODELOPT_NVFP4_BRIDGE_MEMORY_GB,
|
|
"nvfp4_bridge_memory_guard_passed": True,
|
|
"nvfp4_bridge_smoke": {},
|
|
"ready": False,
|
|
"supports_pixels": False,
|
|
"supports_video": False,
|
|
"supports_thinking": False,
|
|
"media_policy": "unknown",
|
|
"notes": [],
|
|
}
|
|
try:
|
|
import importlib.util
|
|
|
|
status["transformers_importable"] = importlib.util.find_spec("transformers") is not None
|
|
status["diffusion_gemma_module"] = importlib.util.find_spec("transformers.models.diffusion_gemma") is not None
|
|
status["vllm_importable"] = importlib.util.find_spec("vllm") is not None
|
|
status["modelopt_importable"] = bool(
|
|
importlib.util.find_spec("modelopt") or importlib.util.find_spec("nvidia_modelopt")
|
|
)
|
|
status["comfy_kitchen_importable"] = importlib.util.find_spec("comfy_kitchen") is not None
|
|
status["fp_quant_importable"] = importlib.util.find_spec("fp_quant") is not None
|
|
status["qutlass_importable"] = importlib.util.find_spec("qutlass") is not None
|
|
status["tensorrt_llm_importable"] = importlib.util.find_spec("tensorrt_llm") is not None
|
|
try:
|
|
import transformers
|
|
|
|
status["transformers_version"] = str(getattr(transformers, "__version__", ""))
|
|
except Exception:
|
|
pass
|
|
try:
|
|
import torch
|
|
|
|
status["torch_cuda_available"] = bool(torch.cuda.is_available())
|
|
if torch.cuda.is_available():
|
|
status["torch_cuda_capability"] = ".".join(str(part) for part in torch.cuda.get_device_capability())
|
|
except Exception as exc:
|
|
if config.backend == "transformers_inprocess":
|
|
status["notes"].append(f"Torch CUDA probe failed: {exc}")
|
|
try:
|
|
from transformers import AutoProcessor # noqa: F401
|
|
|
|
status["auto_processor"] = True
|
|
if config.backend == "transformers_inprocess" and status["model_path_exists"] and status["is_hf_repo"]:
|
|
try:
|
|
processor = AutoProcessor.from_pretrained(config.model_path, local_files_only=True)
|
|
status["processor_class"] = type(processor).__name__
|
|
except Exception as exc:
|
|
status["notes"].append(f"AutoProcessor load failed: {exc}")
|
|
except Exception as exc:
|
|
if config.backend == "transformers_inprocess":
|
|
status["notes"].append(f"AutoProcessor import failed: {exc}")
|
|
try:
|
|
from transformers import DiffusionGemmaForBlockDiffusion # noqa: F401
|
|
|
|
status["diffusion_gemma_class"] = True
|
|
except Exception as exc:
|
|
if config.backend == "transformers_inprocess":
|
|
status["notes"].append(f"DiffusionGemmaForBlockDiffusion import failed: {exc}")
|
|
except Exception as exc:
|
|
if config.backend == "transformers_inprocess":
|
|
status["notes"].append(f"Transformers probe failed: {exc}")
|
|
|
|
if config.backend == "transformers_inprocess":
|
|
modelopt_nvfp4 = status["path_kind"] == "nvfp4_hf_repo" and status["quant_method"] == "modelopt"
|
|
low_bridge_memory = bool(
|
|
modelopt_nvfp4
|
|
and config.quantization == "modelopt_nvfp4"
|
|
and 0 < float(config.max_memory_gb) < MIN_MODELOPT_NVFP4_BRIDGE_MEMORY_GB
|
|
)
|
|
status["nvfp4_bridge_memory_guard_passed"] = not low_bridge_memory
|
|
if modelopt_nvfp4 and config.quantization == "modelopt_nvfp4":
|
|
smoke = _nvfp4_bridge_smoke(config.model_path)
|
|
status["nvfp4_bridge_smoke"] = smoke
|
|
status["comfy_nvfp4_bridge_supported"] = bool(
|
|
smoke.get("passed")
|
|
and status["comfy_kitchen_importable"]
|
|
and status["torch_cuda_available"]
|
|
and status["diffusion_gemma_class"]
|
|
and status["auto_processor"]
|
|
)
|
|
status["modelopt_nvfp4_supported"] = bool(
|
|
modelopt_nvfp4
|
|
and (
|
|
status["comfy_nvfp4_bridge_supported"]
|
|
or (
|
|
status["modelopt_importable"]
|
|
and status["tensorrt_llm_importable"]
|
|
and status["has_modelopt_state"]
|
|
)
|
|
)
|
|
)
|
|
status["ready"] = bool(
|
|
status["model_path_exists"]
|
|
and status["auto_processor"]
|
|
and status["diffusion_gemma_class"]
|
|
and not (modelopt_nvfp4 and not status["modelopt_nvfp4_supported"])
|
|
and not low_bridge_memory
|
|
)
|
|
status["supports_pixels"] = bool(status["ready"])
|
|
status["supports_video"] = bool(status["ready"])
|
|
status["supports_thinking"] = bool(status["ready"])
|
|
status["media_policy"] = "pixels_when_ready"
|
|
if modelopt_nvfp4 and status["comfy_nvfp4_bridge_supported"]:
|
|
status["notes"].append(
|
|
"This NVFP4 repo uses quant_method=modelopt, which vanilla Transformers does not deserialize. "
|
|
"The node will use the local Comfy NVFP4 bridge: Transformers supplies DiffusionGemma structure "
|
|
"and Comfy/comfy_kitchen supplies packed NVFP4 expert matmul."
|
|
)
|
|
if low_bridge_memory:
|
|
status["notes"].append(
|
|
f"The local NVFP4 bridge does not offload by max_memory_gb yet. Set max_memory_gb to "
|
|
f"{MIN_MODELOPT_NVFP4_BRIDGE_MEMORY_GB:g} or higher, or 0 for an explicit full-GPU load. "
|
|
"Lower values are blocked to avoid a partial CUDA load that can require a ComfyUI restart."
|
|
)
|
|
elif modelopt_nvfp4 and not status["modelopt_nvfp4_supported"]:
|
|
status["notes"].append(
|
|
"This NVFP4 repo uses quant_method=modelopt. Transformers 5.12.1 exposes DiffusionGemma, "
|
|
"but it does not register modelopt as an in-process quantizer. The local ModelOpt package "
|
|
"also lacks a modelopt_state.pth restore file for this repo, TensorRT-LLM is not available, "
|
|
"and the Comfy NVFP4 bridge smoke did not pass. "
|
|
"Blocking before full weight load to avoid a CPU-memory hang."
|
|
)
|
|
elif config.backend == "gguf_subprocess":
|
|
status["cli_path_exists"] = bool(config.cli_path and Path(config.cli_path).exists())
|
|
status["cuda_runtime_bin_exists"] = Path(DEFAULT_CUDA_RUNTIME_BIN).exists()
|
|
status["ready"] = bool(status["model_path_exists"] and status["cli_path_exists"] and status["is_gguf_file"])
|
|
status["supports_pixels"] = False
|
|
status["supports_video"] = False
|
|
status["supports_thinking"] = bool(status["ready"])
|
|
status["media_policy"] = "metadata_only"
|
|
if not status["cli_path_exists"]:
|
|
status["notes"].append("GGUF subprocess backend needs a foreground llama-diffusion-cli path.")
|
|
if status["model_path_exists"] and not status["is_gguf_file"]:
|
|
if status["path_kind"] == "nvfp4_hf_repo":
|
|
status["notes"].append(
|
|
"This path is an NVFP4 Hugging Face/vLLM repo, not a .gguf file. "
|
|
"The clean Model Loader runs llama-diffusion-cli and needs the GGUF model path."
|
|
)
|
|
elif status["is_hf_repo"]:
|
|
status["notes"].append(
|
|
"This path is a Hugging Face repo folder, not a .gguf file. "
|
|
"Use Model Loader (Advanced) with transformers_inprocess only after proof gates pass."
|
|
)
|
|
else:
|
|
status["notes"].append("GGUF subprocess backend needs model_path to point directly to a .gguf file.")
|
|
elif config.backend in {"qwen_vl_placeholder", "minicpm_v_placeholder"}:
|
|
status["ready"] = False
|
|
status["notes"].append(
|
|
"Fallback backend lane is reserved for comparative benchmarking; no Qwen/MiniCPM runner is loaded by this pack yet."
|
|
)
|
|
else:
|
|
status["ready"] = config.backend == "template"
|
|
status["supports_thinking"] = False
|
|
status["media_policy"] = "template"
|
|
|
|
if config.backend == "transformers_inprocess" and not status["ready"]:
|
|
if not status["diffusion_gemma_class"]:
|
|
status["notes"].append(
|
|
"This ComfyUI venv can keep the node pack loaded, but it cannot run DiffusionGemma until "
|
|
"a compatible Transformers build exposes DiffusionGemmaForBlockDiffusion."
|
|
)
|
|
if status["path_kind"] == "nvfp4_hf_repo" and not status["comfy_nvfp4_bridge_supported"]:
|
|
status["notes"].append(
|
|
"NVFP4 was detected. Use this path only after a supported in-process ModelOpt NVFP4 loader is "
|
|
"available; the current local foreground fallback remains the GGUF llama-diffusion-cli route."
|
|
)
|
|
return status
|
|
|
|
|
|
def _strip_markdown(text: str) -> str:
|
|
text = re.sub(r"```(?:json|text|markdown)?\s*", "", text, flags=re.IGNORECASE)
|
|
text = text.replace("```", "")
|
|
text = re.sub(r"^\s{0,3}#{1,6}\s+", "", text, flags=re.MULTILINE)
|
|
return text.strip()
|
|
|
|
|
|
def _extract_delimited_final_payload(text: str) -> tuple[str, bool]:
|
|
value = str(text or "")
|
|
if not value:
|
|
return "", False
|
|
lowered = value.lower()
|
|
starts: list[tuple[int, str]] = []
|
|
for marker in FINAL_JSON_OPEN_MARKERS:
|
|
index = lowered.find(marker.lower())
|
|
if index >= 0:
|
|
starts.append((index, marker))
|
|
if not starts:
|
|
return value, False
|
|
start_index, open_marker = min(starts, key=lambda item: item[0])
|
|
payload_start = start_index + len(open_marker)
|
|
payload_end = len(value)
|
|
for marker in FINAL_JSON_CLOSE_MARKERS:
|
|
index = lowered.find(marker.lower(), payload_start)
|
|
if index >= 0:
|
|
payload_end = min(payload_end, index)
|
|
return value[payload_start:payload_end].strip(), True
|
|
|
|
|
|
def _strip_thinking(text: str) -> str:
|
|
text = re.sub(r"<think>.*?</think>", "", text, flags=re.IGNORECASE | re.DOTALL)
|
|
text = re.sub(r"<\|thought\|>.*?<\|end_thought\|>", "", text, flags=re.IGNORECASE | re.DOTALL)
|
|
text = re.sub(r"<\|thought\|>.*?(?=<\|(?:final_json|answer|final)\|>)", "", text, flags=re.IGNORECASE | re.DOTALL)
|
|
text = re.sub(r"<\|channel\|>thought.*?(?=<\|channel\|>|$)", "", text, flags=re.IGNORECASE | re.DOTALL)
|
|
text = re.sub(r"<\|?channel\|?>\s*thought\s*", "", text, flags=re.IGNORECASE)
|
|
text = re.sub(r"<\|?channel\|?>", "", text, flags=re.IGNORECASE)
|
|
text = re.sub(r"<\|[^>]+?\|>", "", text)
|
|
return text.strip()
|
|
|
|
|
|
def _strip_cli_runtime_lines(text: str) -> str:
|
|
kept: list[str] = []
|
|
for line in text.splitlines():
|
|
if re.match(r"\s*(total time:|throughput:|diffusion step:)", line, flags=re.IGNORECASE):
|
|
continue
|
|
if re.match(r"\s*\d+\.\d+\.\d+\.\d+\s+[IWE]\s+", line):
|
|
continue
|
|
kept.append(line)
|
|
return "\n".join(kept).strip()
|
|
|
|
|
|
def _strip_generation_artifacts(text: str) -> str:
|
|
value = str(text or "")
|
|
if not value:
|
|
return ""
|
|
replacements = {
|
|
"\ufeff": "",
|
|
"\ufffd": "",
|
|
"": "",
|
|
"ç´‹": " ",
|
|
}
|
|
for artifact, replacement in replacements.items():
|
|
value = value.replace(artifact, replacement)
|
|
value = re.sub(r"[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f]", " ", value)
|
|
value = re.sub(r"(?:(?<=\s)|(?<=[,.;:!?]))(?:[ÃÂâ][^\sA-Za-z0-9]{1,8}|[ÃÂ][A-Za-z0-9]{1,4})(?=\s|[,.;:!?]|$)", " ", value)
|
|
value = re.sub(r"\s*,\s+to\s+(him|her|them)\b", r" near \1", value, flags=re.IGNORECASE)
|
|
value = re.sub(r"\s+([,.!?;:])", r"\1", value)
|
|
value = re.sub(r"([,.!?;:])(?=[^\s,.!?;:])", r"\1 ", value)
|
|
return re.sub(r"\s+", " ", value).strip()
|
|
|
|
|
|
def _split_extra_args(extra_args: str) -> list[str]:
|
|
if not extra_args.strip():
|
|
return []
|
|
if os.name != "nt":
|
|
return shlex.split(extra_args)
|
|
return [part.strip("\"'") for part in shlex.split(extra_args, posix=False)]
|
|
|
|
|
|
def _format_temperature(value: Any, default: float = 0.2) -> str:
|
|
temp = max(0.0, min(2.0, _safe_float(value, default)))
|
|
return f"{temp:.3f}".rstrip("0").rstrip(".")
|
|
|
|
|
|
def _extra_args_with_temperature(extra_args: str, temperature: float) -> str:
|
|
value = (extra_args or "").strip()
|
|
temp = _format_temperature(temperature)
|
|
if not value:
|
|
return f"--temp {temp}"
|
|
pattern = r"(?<!\S)(?:--temp|-temp|--temperature|-temperature)\s+\S+"
|
|
if re.search(pattern, value):
|
|
return re.sub(pattern, f"--temp {temp}", value, count=1)
|
|
return f"{value} --temp {temp}".strip()
|
|
|
|
|
|
def _subprocess_env_for_cli(cli_path: str) -> dict[str, str]:
|
|
env = os.environ.copy()
|
|
path_key = "Path" if "Path" in env else "PATH"
|
|
existing_path = env.get(path_key, "")
|
|
candidates = [
|
|
str(Path(cli_path).parent) if cli_path else "",
|
|
DEFAULT_CUDA_RUNTIME_BIN,
|
|
]
|
|
prefix = [path for path in candidates if path and Path(path).exists()]
|
|
env[path_key] = os.pathsep.join([*prefix, existing_path]) if existing_path else os.pathsep.join(prefix)
|
|
cuda_root = Path(DEFAULT_CUDA_RUNTIME_BIN).parent
|
|
if cuda_root.exists():
|
|
env.setdefault("CUDA_PATH", str(cuda_root))
|
|
return env
|
|
|
|
|
|
def _sanitize_prompt_text(
|
|
text: str,
|
|
fallback: str = "",
|
|
max_chars: int = 4000,
|
|
strip_thinking: bool = True,
|
|
strip_markdown: bool = True,
|
|
) -> str:
|
|
value = (text or "").strip()
|
|
if strip_thinking:
|
|
value = _strip_thinking(value)
|
|
if strip_markdown:
|
|
value = _strip_markdown(value)
|
|
value = _strip_generation_artifacts(value)
|
|
value = re.sub(r"\s+", " ", value).strip()
|
|
if not value:
|
|
value = (fallback or "").strip()
|
|
if max_chars > 0 and len(value) > max_chars:
|
|
value = value[:max_chars].rsplit(" ", 1)[0].strip() or value[:max_chars].strip()
|
|
return value
|
|
|
|
|
|
def _extract_json_object(text: str) -> dict[str, Any] | None:
|
|
if not text:
|
|
return None
|
|
payload, _had_boundary = _extract_delimited_final_payload(text)
|
|
cleaned = _strip_markdown(_strip_thinking(payload))
|
|
candidates = [cleaned]
|
|
candidates.append(re.sub(r"^\s*\{\s*(?=\{)", "", cleaned).strip())
|
|
first = cleaned.find("{")
|
|
last = cleaned.rfind("}")
|
|
if first >= 0 and last > first:
|
|
candidates.append(cleaned[first : last + 1])
|
|
candidates.append(re.sub(r"^\s*\{\s*(?=\{)", "", cleaned[first : last + 1]).strip())
|
|
for candidate in candidates:
|
|
try:
|
|
data = json.loads(candidate)
|
|
if isinstance(data, dict):
|
|
return data
|
|
except Exception:
|
|
continue
|
|
return None
|
|
|
|
|
|
def _single_paragraph(text: str) -> str:
|
|
return _strip_generation_artifacts(re.sub(r"\s+", " ", (text or "").replace("\n", " ")).strip())
|
|
|
|
|
|
def _ensure_terminal_period(text: str) -> str:
|
|
value = _single_paragraph(text)
|
|
if not value:
|
|
return ""
|
|
if value[-1] in ".!?":
|
|
return value
|
|
return f"{value}."
|
|
|
|
|
|
def _sentence_case(text: str) -> str:
|
|
value = _single_paragraph(text)
|
|
if value and value[0].islower():
|
|
return value[0].upper() + value[1:]
|
|
return value
|
|
|
|
|
|
def _clean_prompt_request_text(text: str) -> str:
|
|
value = _single_paragraph(text)
|
|
if _is_media_reference_only_request(value):
|
|
return ""
|
|
patterns = [
|
|
r"^(?:create|write|generate|make)\s+(?:an?\s+)?(?:ltx(?:[- ]?2(?:\.3)?)?\s+)?(?:script|prompt)\s+for\s+(?:image|video)\s+generation\s+(?:featuring|with|of|about)\s+",
|
|
r"^(?:create|write|generate|make)\s+(?:an?\s+)?(?:ltx(?:[- ]?2(?:\.3)?)?\s+)?(?:script|prompt)\s+(?:featuring|with|of|about)\s+",
|
|
r"^(?:create|write|generate|make)\s+(?:an?\s+)?(?:image|video)\s+(?:generation\s+)?(?:script|prompt)\s+(?:featuring|with|of|about)\s+",
|
|
]
|
|
cleaned = value
|
|
for pattern in patterns:
|
|
cleaned = re.sub(pattern, "", cleaned, flags=re.IGNORECASE).strip()
|
|
cleaned = re.sub(r"\.{2,}", ".", cleaned)
|
|
cleaned = re.sub(r"\s+([,.!?])", r"\1", cleaned).strip()
|
|
return _normalize_hex_text(_sentence_case(cleaned or value))
|
|
|
|
|
|
def _clean_visual_description(text: str, max_chars: int = 3000) -> str:
|
|
value = _strip_generation_artifacts(_strip_markdown(_strip_thinking(str(text or ""))))
|
|
value = re.sub(r"\s+", " ", value).strip()
|
|
if max_chars > 0 and len(value) > max_chars:
|
|
value = value[:max_chars].rsplit(" ", 1)[0].strip() or value[:max_chars].strip()
|
|
return value
|
|
|
|
|
|
def _visual_description_from_metadata(media_metadata: dict[str, Any] | None) -> str:
|
|
if not isinstance(media_metadata, dict):
|
|
return ""
|
|
return _clean_visual_description(str(media_metadata.get("visual_description", "")))
|
|
|
|
|
|
def _normalize_media_synthesis_mode(value: Any) -> str:
|
|
text = str(value or "video_recreation").strip().lower().replace("-", "_").replace(" ", "_")
|
|
if text in {"image_identity_video_control", "identity_image_video_control", "image_identity", "identity_reference"}:
|
|
return "image_identity_video_control"
|
|
return "video_recreation"
|
|
|
|
|
|
def _is_image_identity_video_control(media_metadata: dict[str, Any] | None) -> bool:
|
|
return _normalize_media_synthesis_mode(
|
|
media_metadata.get("media_synthesis_mode") if isinstance(media_metadata, dict) else ""
|
|
) == "image_identity_video_control"
|
|
|
|
|
|
def _synthesis_contract_text(media_metadata: dict[str, Any] | None) -> str:
|
|
if _is_image_identity_video_control(media_metadata):
|
|
return (
|
|
"Use the reference image as the identity and appearance source for the output subject, including face, hair, body type, wardrobe, accessories, styling, and distinguishing visual traits. "
|
|
"Use the video frames only as control structure for pose, depth, canny/edge layout, composition, action, timing, camera motion, blocking, and scene geometry. "
|
|
"Do not copy the control-video subject's clothing, hair, face, body type, accessories, or style unless the user explicitly asks for those video-subject traits."
|
|
)
|
|
return "Recreate the source video or media evidence directly unless the user prompt asks for a different transformation."
|
|
|
|
|
|
def _copy_synthesis_metadata(target: dict[str, Any], media_metadata: dict[str, Any]) -> None:
|
|
for key in (
|
|
"media_synthesis_mode",
|
|
"image_role",
|
|
"video_role",
|
|
"reference_image_count",
|
|
"reference_image_backend_attached",
|
|
"reference_image_resized_for_video_control",
|
|
"reference_image_resize_mode",
|
|
"reference_image_original_width",
|
|
"reference_image_original_height",
|
|
"reference_image_backend_width",
|
|
"reference_image_backend_height",
|
|
"video_sampled_frame_count",
|
|
"synthesis_contract",
|
|
):
|
|
if key in media_metadata:
|
|
target[key] = media_metadata[key]
|
|
|
|
|
|
def _media_source_has_pixels(media_metadata: dict[str, Any] | None) -> bool:
|
|
if not isinstance(media_metadata, dict):
|
|
return False
|
|
source = str(media_metadata.get("source", "none")).lower()
|
|
return source in {"image", "video", "image+video"} or bool(media_metadata.get("sampled_frame_count"))
|
|
|
|
|
|
def _media_pixels_sent_to_backend(media_metadata: dict[str, Any] | None) -> bool:
|
|
return bool(isinstance(media_metadata, dict) and media_metadata.get("pixels_sent_to_backend"))
|
|
|
|
|
|
def _media_grounding_mode(media_metadata: dict[str, Any] | None) -> str:
|
|
if not _media_source_has_pixels(media_metadata):
|
|
return "none"
|
|
if _media_pixels_sent_to_backend(media_metadata):
|
|
return "pixels"
|
|
if _visual_description_from_metadata(media_metadata):
|
|
return "visual_description"
|
|
return "metadata_only"
|
|
|
|
|
|
def _video_transport_mode() -> str:
|
|
value = os.environ.get("DG_VIDEO_TRANSPORT", "sampled_frame_images").strip().lower()
|
|
value = value.replace("-", "_").replace(" ", "_")
|
|
if value in {"video", "video_token", "video_tokens", "native_video"}:
|
|
return "video_tokens"
|
|
return "sampled_frame_images"
|
|
|
|
|
|
def _video_frame_timecodes(media_metadata: dict[str, Any] | None, frame_count: int) -> list[float]:
|
|
if frame_count <= 0:
|
|
return []
|
|
if not isinstance(media_metadata, dict):
|
|
return [float(i) for i in range(frame_count)]
|
|
indices = media_metadata.get("sampled_indices")
|
|
source_fps = _safe_float(media_metadata.get("source_fps"), 0.0)
|
|
if isinstance(indices, list) and source_fps > 0:
|
|
values: list[float] = []
|
|
for index in indices[:frame_count]:
|
|
try:
|
|
values.append(float(index) / source_fps)
|
|
except Exception:
|
|
values.append(float(len(values)))
|
|
if len(values) == frame_count:
|
|
return values
|
|
duration = _safe_float(media_metadata.get("duration_seconds"), 0.0)
|
|
if duration > 0 and frame_count > 1:
|
|
step = duration / max(1, frame_count - 1)
|
|
return [idx * step for idx in range(frame_count)]
|
|
sample_fps = max(0.001, _safe_float(media_metadata.get("sample_fps"), 1.0))
|
|
return [idx / sample_fps for idx in range(frame_count)]
|
|
|
|
|
|
def _display_video_frame_count(media_metadata: dict[str, Any] | None, frame_count: int) -> int:
|
|
if not isinstance(media_metadata, dict):
|
|
return frame_count
|
|
video_count = int(max(0, _safe_float(media_metadata.get("video_sampled_frame_count"), 0.0)))
|
|
if video_count > 0:
|
|
return min(frame_count, video_count)
|
|
return frame_count
|
|
|
|
|
|
def _sampled_video_frame_sequence_intro(media_metadata: dict[str, Any] | None, frame_count: int) -> str:
|
|
reference_count = int(max(0, _safe_float(media_metadata.get("reference_image_count"), 0.0))) if isinstance(media_metadata, dict) else 0
|
|
video_frame_count = _display_video_frame_count(media_metadata, max(0, frame_count - reference_count))
|
|
timecodes = _video_frame_timecodes(media_metadata, video_frame_count)
|
|
shown = ", ".join(f"{idx + 1}={seconds:.2f}s" for idx, seconds in enumerate(timecodes[:24]))
|
|
if video_frame_count > 24:
|
|
shown = f"{shown}, ..."
|
|
source = str(media_metadata.get("source", "video")) if isinstance(media_metadata, dict) else "video"
|
|
if source == "image+video" and _is_image_identity_video_control(media_metadata):
|
|
lead = (
|
|
"The first attached image is the identity reference for the output subject. "
|
|
"The remaining attached images are ordered sampled video control frames."
|
|
)
|
|
return (
|
|
f"{lead} Video frame times: {shown or 'unknown'}. "
|
|
"Describe the reference-image subject performing the video action/control structure. "
|
|
"Use the identity image for face, hair, body type, wardrobe, accessories, styling, and distinguishing visual traits. "
|
|
"Use the video frames for pose, depth, canny/edge layout, blocking, action, timing, composition, camera motion, and scene geometry. "
|
|
"Do not copy the control-video subject's clothing, hair, face, body type, accessories, or styling into the prompt unless the user explicitly asks for those video-subject traits. "
|
|
"Do not let the identity or appearance of the person in the control video override the reference-image subject. "
|
|
"Preserve the opening framing, any visible camera movement or stable hold during the clip, and the closing framing."
|
|
)
|
|
if source == "image+video":
|
|
lead = "The attached images are an optional still reference followed by ordered sampled frames from one video."
|
|
else:
|
|
lead = "The attached images are ordered sampled frames from one video."
|
|
return (
|
|
f"{lead} Frame times: {shown or 'unknown'}. Treat the images as one temporal sequence. "
|
|
"Ground the answer only in visible subjects, clothing, setting, action, camera behavior, and motion that appear in these frames. "
|
|
"Preserve the opening framing, any visible camera movement or stable hold during the clip, and the closing framing. "
|
|
"If the camera is mostly static, say that rather than inventing a zoom, pan, reveal, exit, or wider pull-back. "
|
|
"Do not invent water, animals, vehicles, dialogue, locations, or extra characters unless they are visibly present."
|
|
)
|
|
|
|
|
|
def _media_grounding_instruction(media_metadata: dict[str, Any] | None) -> str:
|
|
mode = _media_grounding_mode(media_metadata)
|
|
visual_description = _visual_description_from_metadata(media_metadata)
|
|
capture_facts = (
|
|
"Also preserve visual capture facts such as exposure level, blown or clipped highlights, flash/direct light, "
|
|
"bloom, halation, lens blur, bokeh, depth of field, focal length impression, perspective compression, lens distortion, "
|
|
"motion blur, grain/noise, white balance, contrast, dynamic range, and camera angle."
|
|
)
|
|
if mode == "pixels":
|
|
source = str(media_metadata.get("source", "none")).lower() if isinstance(media_metadata, dict) else "none"
|
|
transport = str(media_metadata.get("transformers_video_transport", _video_transport_mode())) if isinstance(media_metadata, dict) else _video_transport_mode()
|
|
if source in {"video", "image+video"} and transport == "sampled_frame_images":
|
|
if _is_image_identity_video_control(media_metadata):
|
|
return (
|
|
"Image+video synthesis mode is active. The first attached image is the identity reference for the output subject; "
|
|
"the ordered video frames are the control source for pose, depth, canny/edge layout, blocking, action, timing, composition, camera motion, and scene geometry. "
|
|
"Describe the reference-image subject performing the video action/control structure. The reference image controls face, hair, body type, wardrobe, accessories, styling, and distinguishing visual traits. "
|
|
"Do not copy the control-video subject's clothing, hair, face, body type, accessories, or styling unless explicitly requested, and do not let the video subject's identity dominate the output. "
|
|
"Preserve source-video chronology, opening framing, visible camera movement or stable hold, and closing framing while transferring the identity from the reference image. "
|
|
f"{capture_facts}"
|
|
)
|
|
return (
|
|
"Ordered sampled video frames are attached through the image pathway. Use them as a time sequence, "
|
|
"and base the prompt only on visible facts that appear in those frames. Treat time of day, ambient brightness, "
|
|
"visible light sources, shadows, contrast, and exposure as source-identity facts to preserve. "
|
|
"Preserve the source clip chronologically: opening framing, any visible camera movement or stable hold through the middle, and the closing framing. "
|
|
"If the camera is mostly static, say that. Do not invent zooms, pans, reveals, exits, wider blocking, or lens shifts that are not visibly present. "
|
|
f"{capture_facts}"
|
|
)
|
|
return (
|
|
"Media pixels are attached to the backend request. Use visible facts from those pixels, and do not invent details beyond them. "
|
|
"Treat time of day, ambient brightness, visible light sources, shadows, contrast, and exposure as source-identity facts to preserve. "
|
|
f"{capture_facts}"
|
|
)
|
|
if mode == "visual_description":
|
|
return (
|
|
"Media pixels are not visible to the active backend. Treat this visual_description as the only visual evidence, including any stated time of day, ambient brightness, light sources, shadows, contrast, and exposure: "
|
|
f"{visual_description}. {capture_facts}"
|
|
)
|
|
if mode == "metadata_only":
|
|
return (
|
|
"Image/video tensors were supplied, but the active backend cannot inspect pixels and no visual_description was supplied. "
|
|
"Do not describe the image/video as if seen. Do not use generic phrases such as 'provided scene' or 'high-fidelity production design'. "
|
|
"Use only concrete visual facts present in the user's text; if the user only asks to describe the media, report the limitation in metadata."
|
|
)
|
|
return "No image/video media is attached."
|
|
|
|
|
|
def _is_media_reference_only_request(text: str) -> bool:
|
|
value = _single_paragraph(text).lower()
|
|
if not value:
|
|
return True
|
|
media_words = ("image", "photo", "picture", "video", "clip", "frame", "still", "scene", "media")
|
|
ask_words = ("describe", "analyze", "caption", "interrogate", "what is", "what's", "turn this", "use this", "from this")
|
|
if not any(word in value for word in media_words) or not any(word in value for word in ask_words):
|
|
return False
|
|
descriptive_words = re.findall(r"[a-z0-9][a-z0-9'-]{2,}", value)
|
|
filler = {
|
|
"describe", "analyze", "caption", "interrogate", "what", "this", "that", "image", "photo",
|
|
"picture", "video", "clip", "frame", "still", "scene", "media", "turn", "into", "from",
|
|
"use", "prompt", "ltx", "ideogram", "the", "please", "make", "create", "generate",
|
|
}
|
|
specific = [word for word in descriptive_words if word not in filler]
|
|
return len(specific) <= 2
|
|
|
|
|
|
def _looks_like_unseen_media_bluff(text: str) -> bool:
|
|
value = _single_paragraph(text).lower()
|
|
if not value:
|
|
return False
|
|
phrases = [
|
|
"provided scene",
|
|
"essence of the provided",
|
|
"high-fidelity production design",
|
|
"soft, natural light filters through the atmosphere",
|
|
"subtle, realistic movements",
|
|
"intricate details of their wardrobe",
|
|
"professional film-production aesthetic",
|
|
"balanced color grading",
|
|
"realistic atmospheric depth",
|
|
]
|
|
hits = sum(1 for phrase in phrases if phrase in value)
|
|
return hits >= 2
|
|
|
|
|
|
def _truthy_metadata(value: Any) -> bool:
|
|
if isinstance(value, bool):
|
|
return value
|
|
if isinstance(value, (int, float)):
|
|
return bool(value)
|
|
if isinstance(value, str):
|
|
return value.strip().lower() in {"1", "true", "yes", "y", "on"}
|
|
return False
|
|
|
|
|
|
def _contains_any_phrase(text: str, phrases: tuple[str, ...]) -> bool:
|
|
value = _single_paragraph(text).lower()
|
|
return any(phrase in value for phrase in phrases)
|
|
|
|
|
|
def _has_visual_claim_not_in_user(prompt: str, user_prompt: str, phrases: tuple[str, ...]) -> bool:
|
|
prompt_value = _single_paragraph(prompt).lower()
|
|
user_value = _single_paragraph(user_prompt).lower()
|
|
return any(phrase in prompt_value and phrase not in user_value for phrase in phrases)
|
|
|
|
|
|
def _contains_daylight_claim(text: str) -> bool:
|
|
return _contains_any_phrase(
|
|
text,
|
|
(
|
|
"daylight",
|
|
"broad daylight",
|
|
"sunlit",
|
|
"sun-lit",
|
|
"sunny",
|
|
"blue sky",
|
|
"clear sky",
|
|
"morning light",
|
|
"afternoon light",
|
|
"golden hour",
|
|
),
|
|
)
|
|
|
|
|
|
def _contains_night_evidence(text: str) -> bool:
|
|
return _contains_any_phrase(
|
|
text,
|
|
(
|
|
"night",
|
|
"nighttime",
|
|
"after dark",
|
|
"dark street",
|
|
"low light",
|
|
"moonlit",
|
|
"neon",
|
|
"streetlight",
|
|
),
|
|
)
|
|
|
|
|
|
def _contains_static_camera_evidence(text: str) -> bool:
|
|
return _contains_any_phrase(
|
|
text,
|
|
(
|
|
"static camera",
|
|
"locked off",
|
|
"locked-off",
|
|
"fixed camera",
|
|
"stationary camera",
|
|
"tripod",
|
|
"stable hold",
|
|
"no camera movement",
|
|
"no pan",
|
|
"no zoom",
|
|
),
|
|
)
|
|
|
|
|
|
def _contains_camera_motion_claim(text: str) -> bool:
|
|
return _contains_any_phrase(
|
|
text,
|
|
(
|
|
"camera pans",
|
|
"camera pan",
|
|
"panning",
|
|
"slow pan",
|
|
"camera tilts",
|
|
"tilting",
|
|
"zoom",
|
|
"zooms",
|
|
"push in",
|
|
"pushes in",
|
|
"pull back",
|
|
"pulls back",
|
|
"pull-back",
|
|
"dolly",
|
|
"tracking shot",
|
|
"camera tracks",
|
|
"crane shot",
|
|
"reveal",
|
|
),
|
|
)
|
|
|
|
|
|
def _verified_prompt_claims(
|
|
*,
|
|
prompt_text: str,
|
|
context: GemmaContext,
|
|
media_metadata: dict[str, Any],
|
|
) -> tuple[bool, list[str]]:
|
|
reasons: list[str] = []
|
|
prompt_text = _single_paragraph(_sanitize_prompt_text(prompt_text, "", 8000))
|
|
if not prompt_text:
|
|
reasons.append("empty_generation_prompt")
|
|
return False, reasons
|
|
|
|
grounding_mode = _media_grounding_mode(media_metadata)
|
|
user_prompt = _single_paragraph(_sanitize_prompt_text(context.user_prompt, "", 4000))
|
|
visual_description = _clean_visual_description(
|
|
context.visual_description or str(media_metadata.get("visual_description", ""))
|
|
)
|
|
if grounding_mode == "metadata_only":
|
|
claim_phrases = (
|
|
"daylight",
|
|
"sunlit",
|
|
"sunny",
|
|
"blue sky",
|
|
"night",
|
|
"neon",
|
|
"streetlight",
|
|
"person",
|
|
"character",
|
|
"camera",
|
|
"lens",
|
|
"close-up",
|
|
"wide shot",
|
|
"foreground",
|
|
"background",
|
|
)
|
|
if _is_media_reference_only_request(user_prompt):
|
|
reasons.append("metadata_only_media_reference_request")
|
|
elif _looks_like_unseen_media_bluff(prompt_text):
|
|
reasons.append("metadata_only_unseen_media_bluff")
|
|
elif _has_visual_claim_not_in_user(prompt_text, user_prompt, claim_phrases):
|
|
reasons.append("metadata_only_unverified_visual_claim")
|
|
|
|
if visual_description:
|
|
if _contains_night_evidence(visual_description) and _contains_daylight_claim(prompt_text):
|
|
reasons.append("visual_description_contradiction_daylight_claim")
|
|
if _contains_static_camera_evidence(visual_description) and _contains_camera_motion_claim(prompt_text):
|
|
reasons.append("visual_description_contradiction_camera_motion")
|
|
|
|
return not reasons, reasons
|
|
|
|
|
|
def _packet_generation_readiness(
|
|
packet: dict[str, Any],
|
|
context: GemmaContext,
|
|
target_config: TargetProfileConfig | dict[str, Any],
|
|
prompt_text: str,
|
|
media_metadata: dict[str, Any],
|
|
) -> tuple[bool, list[str]]:
|
|
metadata = packet.get("metadata") if isinstance(packet.get("metadata"), dict) else {}
|
|
reasons: list[str] = []
|
|
json_parse_known = "json_parse_valid" in metadata
|
|
if json_parse_known and not _truthy_metadata(metadata.get("json_parse_valid")):
|
|
reasons.append("json_parse_invalid")
|
|
if _truthy_metadata(metadata.get("plain_text_salvage")) or metadata.get("salvage_warning") or metadata.get("json_parse_warning"):
|
|
reasons.append("salvaged_output")
|
|
if _truthy_metadata(metadata.get("used_template_fallback")) or metadata.get("fallback_reason"):
|
|
reasons.append("template_fallback")
|
|
if _truthy_metadata(metadata.get("template_used_for_missing_fields")):
|
|
reasons.append("template_filled_missing_fields")
|
|
if metadata.get("visual_grounding_warning"):
|
|
reasons.append("metadata_only_visual_grounding_warning")
|
|
if _is_image_identity_video_control(media_metadata):
|
|
if _safe_float(media_metadata.get("reference_image_count"), 0.0) <= 0:
|
|
reasons.append("synthesis_missing_identity_image")
|
|
elif media_metadata.get("reference_image_backend_attached") is False:
|
|
reasons.append("synthesis_identity_image_not_attached")
|
|
if _safe_float(media_metadata.get("video_sampled_frame_count"), 0.0) <= 0:
|
|
reasons.append("synthesis_missing_control_video")
|
|
|
|
claims_ok, claim_reasons = _verified_prompt_claims(
|
|
prompt_text=prompt_text,
|
|
context=context,
|
|
media_metadata=media_metadata,
|
|
)
|
|
if not claims_ok:
|
|
reasons.extend(claim_reasons)
|
|
|
|
deduped: list[str] = []
|
|
seen: set[str] = set()
|
|
for reason in reasons:
|
|
clean = str(reason or "").strip()
|
|
if clean and clean not in seen:
|
|
seen.add(clean)
|
|
deduped.append(clean)
|
|
return not deduped, deduped
|
|
|
|
|
|
def _split_segments(prompt: str, duration_seconds: float, max_segments: int = 6) -> list[dict[str, Any]]:
|
|
protected = re.sub(r"\b(EXT|INT)\.", r"\1<DOT>", prompt or "", flags=re.IGNORECASE)
|
|
sentences = [s.strip().replace("<DOT>", ".") for s in re.split(r"(?<=[.!?])\s+", protected) if s.strip()]
|
|
if len(sentences) > 1 and sentences[0].lower().startswith("style:"):
|
|
sentences = [f"{sentences[0]} {sentences[1]}".strip(), *sentences[2:]]
|
|
if not sentences:
|
|
sentences = [prompt.strip()] if prompt.strip() else ["A clear scene with a visible subject and motion."]
|
|
max_segments = max(1, min(max_segments, len(sentences)))
|
|
chunks: list[str] = []
|
|
for idx in range(max_segments):
|
|
chunks.append(sentences[idx])
|
|
if len(sentences) > max_segments:
|
|
chunks[-1] = " ".join([chunks[-1], *sentences[max_segments:]]).strip()
|
|
segment_duration = duration_seconds / len(chunks) if duration_seconds > 0 else 0.0
|
|
return [
|
|
{
|
|
"index": idx,
|
|
"duration_seconds": round(segment_duration, 3) if segment_duration else 0.0,
|
|
"prompt": chunk,
|
|
}
|
|
for idx, chunk in enumerate(chunks)
|
|
]
|
|
|
|
|
|
def _segments_to_director_strings(segments: list[dict[str, Any]]) -> tuple[str, str]:
|
|
prompts = [str(seg.get("prompt", "")).strip() for seg in segments if str(seg.get("prompt", "")).strip()]
|
|
lengths = []
|
|
for seg in segments[: len(prompts)]:
|
|
duration = _safe_float(seg.get("duration_seconds"), 0.0)
|
|
lengths.append(str(max(1, int(round(duration)))) if duration > 0 else "1")
|
|
return " | ".join(prompts), ", ".join(lengths)
|
|
|
|
|
|
def _load_scene_segments(scene_segments_json: str) -> list[dict[str, Any]] | None:
|
|
parsed = _extract_json_object(scene_segments_json)
|
|
if parsed and isinstance(parsed.get("scene_segments"), list):
|
|
return parsed["scene_segments"]
|
|
if scene_segments_json.strip().startswith("["):
|
|
try:
|
|
loaded = json.loads(scene_segments_json)
|
|
if isinstance(loaded, list):
|
|
return loaded
|
|
except Exception:
|
|
return None
|
|
return None
|
|
|
|
|
|
def _normalize_segments(
|
|
segments: list[dict[str, Any]] | None,
|
|
fallback_prompt: str,
|
|
duration_seconds: float,
|
|
max_segments: int,
|
|
) -> list[dict[str, Any]]:
|
|
limit = max(1, min(int(max_segments), 24))
|
|
cleaned: list[dict[str, Any]] = []
|
|
for segment in segments or []:
|
|
if isinstance(segment, dict):
|
|
prompt = _sanitize_prompt_text(str(segment.get("prompt", "")), "", 3000)
|
|
duration = _safe_float(segment.get("duration_seconds"), 0.0)
|
|
else:
|
|
prompt = _sanitize_prompt_text(str(segment), "", 3000)
|
|
duration = 0.0
|
|
if prompt:
|
|
cleaned.append({"duration_seconds": max(0.0, duration), "prompt": prompt})
|
|
|
|
if not cleaned:
|
|
return _split_segments(fallback_prompt, duration_seconds, limit)
|
|
|
|
if len(cleaned) > limit:
|
|
head = cleaned[: limit - 1]
|
|
tail = cleaned[limit - 1 :]
|
|
merged_prompt = _single_paragraph(" ".join(str(item.get("prompt", "")) for item in tail))
|
|
merged_duration = sum(_safe_float(item.get("duration_seconds"), 0.0) for item in tail)
|
|
cleaned = [*head, {"duration_seconds": merged_duration, "prompt": merged_prompt}]
|
|
|
|
if duration_seconds > 0:
|
|
segment_duration = duration_seconds / len(cleaned)
|
|
for item in cleaned:
|
|
item["duration_seconds"] = round(segment_duration, 3)
|
|
|
|
return [
|
|
{
|
|
"index": idx,
|
|
"duration_seconds": round(_safe_float(item.get("duration_seconds"), 0.0), 3),
|
|
"prompt": _single_paragraph(str(item.get("prompt", ""))),
|
|
}
|
|
for idx, item in enumerate(cleaned)
|
|
]
|
|
|
|
|
|
def _normalize_audio_mode(audio_mode: str) -> str:
|
|
value = (audio_mode or "").strip().lower()
|
|
if value in {"auto_scene_audio", "explicit_sound_design", "visual_only"}:
|
|
return value
|
|
return "auto_scene_audio"
|
|
|
|
|
|
def _audio_guidance_text(audio_guidance: str) -> str:
|
|
return _sanitize_prompt_text(audio_guidance, "", 700, strip_thinking=True, strip_markdown=True).rstrip(".")
|
|
|
|
|
|
def _contains_audio_terms(prompt: str) -> bool:
|
|
return bool(
|
|
re.search(
|
|
r"\b(audio|sound|sounds|sonic|music|voice|voices|dialogue|ambience|ambient|foley|silence|silent|"
|
|
r"croak|croaks|chirp|chirps|rustle|rustles|rustling|footstep|footsteps|whisper|whispers|hum|hums|humming|"
|
|
r"buzz|buzzes|buzzing|ring|rings|ringing|rain|rains|raining|patter|patters|thunder|waves|traffic|engine|engines)\b",
|
|
prompt or "",
|
|
flags=re.IGNORECASE,
|
|
)
|
|
)
|
|
|
|
|
|
def _audio_instruction(audio_mode: str, audio_guidance: str) -> str:
|
|
mode = _normalize_audio_mode(audio_mode)
|
|
guidance = _audio_guidance_text(audio_guidance)
|
|
if mode == "visual_only":
|
|
return "Do not include sound, music, dialogue, ambience, or other audio cues in the LTX prompt."
|
|
if guidance:
|
|
return (
|
|
"Weave this sound design into the LTX prompt as part of the scene, not as a generic appended sentence: "
|
|
f"{guidance}."
|
|
)
|
|
if mode == "explicit_sound_design":
|
|
return (
|
|
"Invent concrete diegetic sound design from the visible action and setting, and weave it into the LTX prompt. "
|
|
"Mention specific sources such as movement, environment, voices, machinery, weather, room tone, or silence only when they fit."
|
|
)
|
|
return (
|
|
"When useful for video, weave one concise, scene-specific audio cue into the LTX prompt itself. "
|
|
"Avoid generic lines like natural scene audio supports the motion."
|
|
)
|
|
|
|
|
|
def _clean_control_text(text: str, max_chars: int = 500) -> str:
|
|
return _sanitize_prompt_text(text, "", max_chars, strip_thinking=True, strip_markdown=True).rstrip(".")
|
|
|
|
|
|
def _effective_duration_seconds(media_metadata: dict[str, Any], target_duration_seconds: float) -> float:
|
|
requested = _safe_float(target_duration_seconds, 0.0)
|
|
if requested > 0:
|
|
return requested
|
|
return _safe_float(media_metadata.get("duration_seconds"), 0.0)
|
|
|
|
|
|
def _target_context(
|
|
media_metadata: dict[str, Any],
|
|
target_duration_seconds: float,
|
|
ltx_style: str,
|
|
ideogram_aspect_ratio: str,
|
|
ideogram_render_style: str,
|
|
ideogram_exact_text: str,
|
|
ideogram_json_output: bool,
|
|
creativity_mode: str = "editorial",
|
|
creative_strength: float = 0.6,
|
|
negative_prompt_mode: str = "auto",
|
|
negative_prompt_guidance: str = "",
|
|
) -> dict[str, Any]:
|
|
return {
|
|
"ltx_duration_seconds": round(_effective_duration_seconds(media_metadata, target_duration_seconds), 3),
|
|
"ltx_style": _clean_control_text(ltx_style, 240),
|
|
"ideogram_aspect_ratio": ideogram_aspect_ratio,
|
|
"ideogram_render_style": _clean_control_text(ideogram_render_style, 240),
|
|
"ideogram_exact_text": _sanitize_prompt_text(
|
|
ideogram_exact_text,
|
|
"",
|
|
1000,
|
|
strip_thinking=True,
|
|
strip_markdown=True,
|
|
),
|
|
"ideogram_json_output": bool(ideogram_json_output),
|
|
"creativity_mode": _normalize_creativity_mode(creativity_mode),
|
|
"creative_strength": _normalize_creative_strength(creative_strength),
|
|
"negative_prompt_mode": _normalize_negative_prompt_mode(negative_prompt_mode),
|
|
"negative_prompt_guidance": _negative_prompt_guidance_text(negative_prompt_guidance),
|
|
}
|
|
|
|
|
|
def _normalize_creativity_mode(creativity_mode: str) -> str:
|
|
value = (creativity_mode or "").strip().lower().replace("-", "_").replace(" ", "_")
|
|
if value in {"faithful", "editorial", "cinematic", "concept_art", "wild"}:
|
|
return value
|
|
return "editorial"
|
|
|
|
|
|
def _normalize_creative_strength(creative_strength: Any) -> float:
|
|
return round(max(0.0, min(1.5, _safe_float(creative_strength, 0.6))), 3)
|
|
|
|
|
|
def _creativity_instruction(creativity_mode: str, creative_strength: float) -> str:
|
|
mode = _normalize_creativity_mode(creativity_mode)
|
|
strength = _normalize_creative_strength(creative_strength)
|
|
if strength <= 0.05 or mode == "faithful":
|
|
return (
|
|
"Faithful mode: preserve the user prompt and supplied media tightly. Add only necessary concrete details for clarity; "
|
|
"do not add new props, eras, wardrobe, materials, worldbuilding, or stylistic surprises unless requested."
|
|
)
|
|
if mode == "editorial":
|
|
return (
|
|
f"Editorial creativity strength {strength}: preserve the core request while adding tasteful art direction: "
|
|
"specific materials, lighting, palette, spatial layout, styling, typography, and product or character details."
|
|
)
|
|
if mode == "cinematic":
|
|
return (
|
|
f"Cinematic creativity strength {strength}: preserve the core request while adding filmic production design: "
|
|
"era cues, set design, lens/camera language, blocking, lighting contrast, atmosphere, props, wardrobe, and sound when relevant."
|
|
)
|
|
if mode == "concept_art":
|
|
return (
|
|
f"Concept-art creativity strength {strength}: preserve the subject and exact requested text while adding distinctive worldbuilding, "
|
|
"genre, environment, props, silhouettes, materials, color motifs, and memorable visual hooks."
|
|
)
|
|
return (
|
|
f"Wild creativity strength {strength}: preserve the required subject, exact text, and safety constraints, but push bolder art direction, "
|
|
"surprising genre fusion, unusual set details, dramatic composition, and memorable specifics. Keep the output coherent and usable."
|
|
)
|
|
|
|
|
|
def _target_structure_anchor(target_profile: str, ideogram_json_output: bool) -> str:
|
|
profile = _normalize_target_profile(target_profile)
|
|
if profile == "ideogram4":
|
|
if ideogram_json_output:
|
|
return (
|
|
"Ideogram JSON: source time-of-day, exposure, lens/depth-of-field, and lighting -> background -> primary subject/elements -> spatial lighting/style -> exact quoted text content. "
|
|
"Map these into the caption JSON schema, not prose instructions."
|
|
)
|
|
return (
|
|
"Ideogram prompt: source time-of-day, exposure, lens/depth-of-field, and lighting -> background -> primary subject -> spatial lighting/style -> exact quoted text content. "
|
|
"Keep text placement and typography concrete."
|
|
)
|
|
return (
|
|
"LTX prompt: source identity -> first-frame framing -> chronological subject action -> observed camera choreography and camera-subject relationship -> "
|
|
"lens/focus/exposure/blocking -> lighting/style -> ending composition -> sound when enabled. Write one continuous ready-to-generate paragraph."
|
|
)
|
|
|
|
|
|
def _normalize_negative_prompt_mode(mode: str) -> str:
|
|
value = (mode or "").strip().lower().replace("-", "_").replace(" ", "_")
|
|
if value in {"auto", "empty", "custom"}:
|
|
return value
|
|
return "auto"
|
|
|
|
|
|
def _negative_prompt_guidance_text(text: str) -> str:
|
|
return _sanitize_prompt_text(text, "", 1200, strip_thinking=True, strip_markdown=True).rstrip(".")
|
|
|
|
|
|
def _negative_prompt_instruction(mode: str, guidance: str) -> str:
|
|
mode = _normalize_negative_prompt_mode(mode)
|
|
guidance = _negative_prompt_guidance_text(guidance)
|
|
if mode == "empty":
|
|
return "Return an empty negative_prompt string."
|
|
if mode == "custom":
|
|
if guidance:
|
|
return f"Use this exact negative_prompt string and do not add to it: {guidance}"
|
|
return "Return an empty negative_prompt string because custom mode has no supplied text."
|
|
return (
|
|
"Write a short comma-separated negative_prompt only when it will help the target generator avoid obvious artifacts, "
|
|
"unwanted text, extra subjects, bad anatomy, or layout failures. Return an empty string when no negative prompt is needed."
|
|
)
|
|
|
|
|
|
def _apply_negative_prompt_policy(value: Any, mode: str, guidance: str) -> str:
|
|
mode = _normalize_negative_prompt_mode(mode)
|
|
guidance = _negative_prompt_guidance_text(guidance)
|
|
if mode == "empty":
|
|
return ""
|
|
if mode == "custom":
|
|
return guidance
|
|
return _sanitize_prompt_text(str(value or ""), "", 1200)
|
|
|
|
|
|
def _creative_art_direction_phrase(creativity_mode: str, creative_strength: float, target: str = "image") -> str:
|
|
mode = _normalize_creativity_mode(creativity_mode)
|
|
strength = _normalize_creative_strength(creative_strength)
|
|
if strength <= 0.05 or mode == "faithful":
|
|
return ""
|
|
if mode == "editorial":
|
|
return "Refined editorial art direction adds premium materials, deliberate styling, polished lighting, and a clear visual hierarchy."
|
|
if mode == "cinematic":
|
|
return "Cinematic art direction adds film-set texture, purposeful blocking, lens-aware composition, motivated lighting, and atmospheric detail."
|
|
if mode == "concept_art":
|
|
return "Concept-art direction adds distinctive worldbuilding, era cues, unusual props, bold silhouettes, and a memorable visual hook."
|
|
return "Bold surreal art direction adds unexpected genre fusion, striking set details, dramatic composition, and vivid memorable specifics while keeping the core subject recognizable."
|
|
|
|
|
|
def _normalize_target_profile(target_profile: str) -> str:
|
|
value = (target_profile or "").strip().lower().replace("-", "_").replace(".", "_")
|
|
if value in {"ltx2", "ltx2_3", "ltx_2_3", "ltx"}:
|
|
return "ltx"
|
|
if value in {"ideogram", "ideogram4", "ideogram4_json", "ideogram_4"}:
|
|
return "ideogram4"
|
|
if value == "both":
|
|
return "ltx"
|
|
return "ltx"
|
|
|
|
|
|
def _normalize_hex_palette(*texts: str, max_colors: int = 16) -> list[str]:
|
|
colors: list[str] = []
|
|
for text in texts:
|
|
for match in re.findall(r"#(?:[0-9a-fA-F]{3}|[0-9a-fA-F]{6})\b", text or ""):
|
|
raw = match.upper()
|
|
if len(raw) == 4:
|
|
raw = "#" + "".join(ch * 2 for ch in raw[1:])
|
|
if raw not in colors:
|
|
colors.append(raw)
|
|
if len(colors) >= max_colors:
|
|
return colors
|
|
return colors
|
|
|
|
|
|
def _normalize_hex_text(text: str) -> str:
|
|
def repl(match: re.Match[str]) -> str:
|
|
raw = match.group(0).upper()
|
|
if len(raw) == 4:
|
|
return "#" + "".join(ch * 2 for ch in raw[1:])
|
|
return raw
|
|
|
|
return re.sub(r"#(?:[0-9a-fA-F]{3}|[0-9a-fA-F]{6})\b", repl, text or "")
|
|
|
|
|
|
def _is_photo_style(base_prompt: str, render_style: str, style_description: dict[str, Any] | None = None) -> bool:
|
|
text = f"{base_prompt} {render_style}"
|
|
if re.search(r"\b(photo|photograph|photoreal|photorealistic|camera|lens|cinematic|portrait|macro|product photography)\b", text, re.IGNORECASE):
|
|
return True
|
|
if re.search(r"\b(illustration|poster|typography|logo|vector|painting|anime|comic|3d render|graphic design|isometric)\b", text, re.IGNORECASE):
|
|
return False
|
|
if isinstance(style_description, dict):
|
|
if "photo" in style_description or str(style_description.get("medium", "")).lower() == "photograph":
|
|
return True
|
|
if "art_style" in style_description:
|
|
return False
|
|
return False
|
|
|
|
|
|
def _coerce_string(value: Any, fallback: str = "") -> str:
|
|
text = _single_paragraph(str(value or ""))
|
|
return _normalize_hex_text(text or fallback)
|
|
|
|
|
|
def _ideogram_style_description(render_style: str, exact_text: str) -> dict[str, Any]:
|
|
style = _clean_control_text(render_style, 240)
|
|
is_photo = _is_photo_style("", style)
|
|
aesthetics = style or ("clean, legible, balanced" if exact_text.strip() else "natural, coherent")
|
|
palette = _normalize_hex_palette(style, exact_text, max_colors=16)
|
|
if is_photo:
|
|
result = {
|
|
"aesthetics": aesthetics,
|
|
"lighting": "clear, controlled lighting that supports the requested subject",
|
|
"photo": "realistic photographic detail with natural perspective",
|
|
"medium": "photograph",
|
|
}
|
|
if palette:
|
|
result["color_palette"] = palette
|
|
return result
|
|
result = {
|
|
"aesthetics": aesthetics,
|
|
"lighting": "clear, balanced lighting",
|
|
"medium": "graphic_design" if exact_text.strip() else "illustration",
|
|
"art_style": style or "clean, coherent visual design",
|
|
}
|
|
if palette:
|
|
result["color_palette"] = palette
|
|
return result
|
|
|
|
|
|
def _ideogram_prompt_object(
|
|
base_prompt: str,
|
|
aspect_ratio: str,
|
|
render_style: str,
|
|
exact_text: str,
|
|
creativity_mode: str = "editorial",
|
|
creative_strength: float = 0.6,
|
|
) -> dict[str, Any]:
|
|
base = _normalize_hex_text(_clean_prompt_request_text(base_prompt) or "A clean, coherent image matching the user's request.")
|
|
art_direction = _creative_art_direction_phrase(creativity_mode, creative_strength, "image")
|
|
element_desc = f"{base} Composition is planned for a {aspect_ratio} frame."
|
|
if art_direction:
|
|
element_desc = f"{element_desc} {art_direction}"
|
|
background = "supporting background that does not compete with the main subject"
|
|
if art_direction:
|
|
background = f"Background supports the requested subject with cohesive environment design. {art_direction}"
|
|
elements: list[dict[str, Any]] = [
|
|
{
|
|
"type": "obj",
|
|
"desc": element_desc,
|
|
}
|
|
]
|
|
if exact_text.strip():
|
|
elements.append(
|
|
{
|
|
"type": "text",
|
|
"text": exact_text.strip(),
|
|
"desc": "Exact readable typography integrated into the requested layout.",
|
|
}
|
|
)
|
|
return {
|
|
"high_level_description": base,
|
|
"style_description": _ideogram_style_description(render_style, exact_text),
|
|
"compositional_deconstruction": {
|
|
"background": background,
|
|
"elements": elements,
|
|
},
|
|
}
|
|
|
|
|
|
def _normalize_bbox(value: Any) -> list[int] | None:
|
|
if not isinstance(value, (list, tuple)) or len(value) != 4:
|
|
return None
|
|
bbox: list[int] = []
|
|
for item in value:
|
|
try:
|
|
bbox.append(max(0, min(1000, int(round(float(item))))))
|
|
except Exception:
|
|
return None
|
|
return bbox
|
|
|
|
|
|
def _normalize_ideogram_element(element: Any, exact_text: str = "") -> dict[str, Any] | None:
|
|
if not isinstance(element, dict):
|
|
desc = _single_paragraph(str(element or ""))
|
|
return {"type": "obj", "desc": desc} if desc else None
|
|
|
|
element_type = str(element.get("type", "")).strip().lower()
|
|
text_value = _single_paragraph(str(element.get("text", "") or ""))
|
|
if element_type not in {"obj", "text"}:
|
|
element_type = "text" if text_value else "obj"
|
|
|
|
desc = _coerce_string(
|
|
element.get("desc")
|
|
or element.get("description")
|
|
or element.get("placement")
|
|
or element.get("layout")
|
|
or element.get("prompt"),
|
|
"Main requested visual element.",
|
|
)
|
|
bbox = _normalize_bbox(element.get("bbox"))
|
|
palette = _normalize_hex_palette(
|
|
" ".join(str(item) for item in element.get("color_palette", []) if isinstance(element.get("color_palette"), list)),
|
|
desc,
|
|
max_colors=5,
|
|
)
|
|
|
|
result: dict[str, Any] = {"type": element_type}
|
|
if bbox:
|
|
result["bbox"] = bbox
|
|
if element_type == "text":
|
|
result["text"] = text_value or exact_text.strip()
|
|
result["desc"] = desc
|
|
if palette:
|
|
result["color_palette"] = palette
|
|
if element_type == "text" and not result.get("text"):
|
|
return None
|
|
return result
|
|
|
|
|
|
def _repair_ideogram_caption_object(
|
|
value: Any,
|
|
fallback_prompt: str,
|
|
aspect_ratio: str,
|
|
render_style: str,
|
|
exact_text: str,
|
|
creativity_mode: str = "editorial",
|
|
creative_strength: float = 0.6,
|
|
) -> dict[str, Any]:
|
|
source = value if isinstance(value, dict) else {}
|
|
fallback = _ideogram_prompt_object(fallback_prompt, aspect_ratio, render_style, exact_text, creativity_mode, creative_strength)
|
|
high_level = _coerce_string(source.get("high_level_description"), fallback["high_level_description"])
|
|
|
|
source_style = source.get("style_description") if isinstance(source.get("style_description"), dict) else {}
|
|
is_photo = _is_photo_style(high_level, render_style, source_style)
|
|
palette = _normalize_hex_palette(
|
|
" ".join(str(item) for item in source_style.get("color_palette", []) if isinstance(source_style.get("color_palette"), list)),
|
|
render_style,
|
|
high_level,
|
|
exact_text,
|
|
max_colors=16,
|
|
)
|
|
aesthetics = _coerce_string(source_style.get("aesthetics"), _ideogram_style_description(render_style, exact_text).get("aesthetics", "clean, coherent"))
|
|
lighting = _coerce_string(source_style.get("lighting"), _ideogram_style_description(render_style, exact_text).get("lighting", "clear, balanced lighting"))
|
|
if is_photo:
|
|
style_description = {
|
|
"aesthetics": aesthetics,
|
|
"lighting": lighting,
|
|
"photo": _coerce_string(source_style.get("photo") or render_style, "realistic photographic detail with natural perspective"),
|
|
"medium": "photograph",
|
|
}
|
|
else:
|
|
style_description = {
|
|
"aesthetics": aesthetics,
|
|
"lighting": lighting,
|
|
"medium": _coerce_string(source_style.get("medium"), "graphic_design" if exact_text.strip() else "illustration"),
|
|
"art_style": _coerce_string(source_style.get("art_style") or render_style, "clean, coherent visual design"),
|
|
}
|
|
if palette:
|
|
style_description["color_palette"] = palette
|
|
|
|
source_comp = source.get("compositional_deconstruction") if isinstance(source.get("compositional_deconstruction"), dict) else {}
|
|
background = _coerce_string(
|
|
source_comp.get("background"),
|
|
fallback["compositional_deconstruction"]["background"],
|
|
)
|
|
elements: list[dict[str, Any]] = []
|
|
for element in source_comp.get("elements", []) if isinstance(source_comp.get("elements"), list) else []:
|
|
normalized = _normalize_ideogram_element(element, exact_text)
|
|
if normalized:
|
|
elements.append(normalized)
|
|
if not elements:
|
|
elements = fallback["compositional_deconstruction"]["elements"]
|
|
if exact_text.strip() and not any(item.get("type") == "text" and item.get("text") == exact_text.strip() for item in elements):
|
|
elements.append({"type": "text", "text": exact_text.strip(), "desc": "Exact readable typography integrated into the requested layout."})
|
|
|
|
return {
|
|
"high_level_description": high_level,
|
|
"style_description": style_description,
|
|
"compositional_deconstruction": {
|
|
"background": background,
|
|
"elements": elements,
|
|
},
|
|
}
|
|
|
|
|
|
def _ideogram_prompt_value_to_text(
|
|
value: Any,
|
|
fallback_prompt: str,
|
|
aspect_ratio: str,
|
|
render_style: str,
|
|
exact_text: str,
|
|
json_output: bool,
|
|
max_chars: int,
|
|
creativity_mode: str = "editorial",
|
|
creative_strength: float = 0.6,
|
|
) -> str:
|
|
if json_output:
|
|
candidate = value
|
|
if isinstance(candidate, str):
|
|
parsed = _extract_json_object(candidate)
|
|
candidate = parsed if parsed is not None else {}
|
|
repaired = _repair_ideogram_caption_object(
|
|
candidate,
|
|
fallback_prompt,
|
|
aspect_ratio,
|
|
render_style,
|
|
exact_text,
|
|
creativity_mode,
|
|
creative_strength,
|
|
)
|
|
text = _ideogram_json_dumps(repaired)
|
|
return _strip_thinking(text).strip()
|
|
return _prompt_field_to_text(
|
|
value if not isinstance(value, dict) else _ideogram_prompt_text(fallback_prompt, aspect_ratio, render_style, exact_text),
|
|
_ideogram_prompt_text(fallback_prompt, aspect_ratio, render_style, exact_text),
|
|
max_chars,
|
|
)
|
|
|
|
|
|
def _ideogram_prompt_text(
|
|
base_prompt: str,
|
|
aspect_ratio: str,
|
|
render_style: str,
|
|
exact_text: str,
|
|
) -> str:
|
|
clauses = [_single_paragraph(base_prompt) or "A clean, coherent image matching the user's request."]
|
|
clauses.append(f"Aspect ratio {aspect_ratio}.")
|
|
style = _clean_control_text(render_style, 240)
|
|
if style:
|
|
clauses.append(f"Style: {style}.")
|
|
if exact_text.strip():
|
|
clauses.append(f"Render the exact readable text: {exact_text.strip()}.")
|
|
clauses.append("Use deliberate layout, clear hierarchy, readable typography, and clean negative space.")
|
|
return _single_paragraph(" ".join(clauses))
|
|
|
|
|
|
def _prompt_field_to_text(value: Any, fallback: str, max_chars: int) -> str:
|
|
if isinstance(value, (dict, list)):
|
|
text = _prompt_json_dumps(value)
|
|
else:
|
|
text = str(value or "")
|
|
return _sanitize_prompt_text(text, fallback, max_chars)
|
|
|
|
|
|
def _apply_ltx2_controls(
|
|
prompt: str,
|
|
duration_seconds: float,
|
|
style: str,
|
|
include_audio: bool,
|
|
audio_guidance: str = "",
|
|
) -> str:
|
|
value = _single_paragraph(prompt)
|
|
if not value:
|
|
return ""
|
|
clauses: list[str] = []
|
|
style_value = _sanitize_prompt_text(style, "", 300, strip_thinking=True, strip_markdown=True).rstrip(".")
|
|
if style_value:
|
|
clauses.append(f"Style: {style_value}.")
|
|
clauses.append(value)
|
|
combined = _single_paragraph(" ".join(clauses))
|
|
if include_audio:
|
|
guidance = _audio_guidance_text(audio_guidance)
|
|
if guidance and not _contains_audio_terms(combined):
|
|
combined = f"{combined} Audio: {guidance}."
|
|
return _single_paragraph(combined)
|
|
|
|
|
|
def _metadata_from_media(media_context: MediaContext | None) -> dict[str, Any]:
|
|
if not isinstance(media_context, MediaContext):
|
|
return {"source": "none"}
|
|
return media_context.metadata.copy()
|
|
|
|
|
|
def _media_context_from_gemma_context(context: GemmaContext | None) -> MediaContext | None:
|
|
if not isinstance(context, GemmaContext):
|
|
return None
|
|
if context.images is None:
|
|
return None
|
|
return MediaContext(images=context.images, source=context.source, metadata=context.media_metadata.copy())
|
|
|
|
|
|
def _gemma_context_from_media(
|
|
user_prompt: str,
|
|
media_context: MediaContext | None = None,
|
|
visual_description: str = "",
|
|
) -> GemmaContext:
|
|
metadata = _metadata_from_media(media_context)
|
|
cleaned_visual_description = _clean_visual_description(visual_description or str(metadata.get("visual_description", "")))
|
|
if cleaned_visual_description:
|
|
metadata["visual_description"] = cleaned_visual_description
|
|
metadata.setdefault("visual_description_source", "upstream_text")
|
|
warnings = list(metadata.get("warnings", [])) if isinstance(metadata.get("warnings"), list) else []
|
|
return GemmaContext(
|
|
user_prompt=str(user_prompt or ""),
|
|
images=media_context.images if isinstance(media_context, MediaContext) else None,
|
|
source=str(metadata.get("source", "none")),
|
|
media_metadata=metadata,
|
|
visual_description=cleaned_visual_description,
|
|
warnings=warnings,
|
|
)
|
|
|
|
|
|
def _gemma_context_to_jsonable(context: GemmaContext | None) -> dict[str, Any]:
|
|
if not isinstance(context, GemmaContext):
|
|
return {"user_prompt": "", "source": "none", "media": {"source": "none"}, "warnings": []}
|
|
metadata = context.media_metadata.copy()
|
|
metadata["pixel_tensor_present"] = bool(context.images is not None)
|
|
return {
|
|
"user_prompt": _sanitize_prompt_text(context.user_prompt, "", 2000),
|
|
"source": context.source,
|
|
"visual_description": _clean_visual_description(context.visual_description),
|
|
"media": metadata,
|
|
"warnings": list(context.warnings),
|
|
}
|
|
|
|
|
|
def _gemma_context_preview(context: GemmaContext | None) -> str:
|
|
data = _gemma_context_to_jsonable(context)
|
|
media = data.get("media", {}) if isinstance(data.get("media"), dict) else {}
|
|
bits = [
|
|
f"Prompt: {_sanitize_prompt_text(str(data.get('user_prompt', '')), '', 300)}",
|
|
f"Media source: {media.get('source', data.get('source', 'none'))}",
|
|
]
|
|
if media.get("duration_seconds"):
|
|
bits.append(f"Duration: {float(media.get('duration_seconds')):.2f}s")
|
|
if media.get("sampled_frame_count"):
|
|
bits.append(f"Sampled frames: {media.get('sampled_frame_count')}")
|
|
visual_description = _clean_visual_description(str(data.get("visual_description", "")))
|
|
if visual_description:
|
|
bits.append(f"Visual description: {_sanitize_prompt_text(visual_description, '', 500)}")
|
|
warnings = data.get("warnings", [])
|
|
if warnings:
|
|
bits.append("Warnings: " + "; ".join(str(item) for item in warnings[:4]))
|
|
return "\n".join(bits)
|
|
|
|
|
|
def _target_profile_config_to_dict(config: TargetProfileConfig | dict[str, Any] | None) -> dict[str, Any]:
|
|
if isinstance(config, TargetProfileConfig):
|
|
data = asdict(config)
|
|
elif isinstance(config, dict):
|
|
data = config.copy()
|
|
else:
|
|
data = {}
|
|
return {
|
|
"target_profile": _normalize_target_profile(str(data.get("target_profile", "ltx"))),
|
|
"audio_mode": _normalize_audio_mode(str(data.get("audio_mode", "auto_scene_audio"))),
|
|
"audio_guidance": _audio_guidance_text(str(data.get("audio_guidance", ""))),
|
|
"target_duration_seconds": max(0.0, _safe_float(data.get("target_duration_seconds"), 0.0)),
|
|
"ltx_style": _clean_control_text(str(data.get("ltx_style", "")), 240),
|
|
"ideogram_aspect_ratio": str(data.get("ideogram_aspect_ratio", "1:1") or "1:1"),
|
|
"ideogram_render_style": _clean_control_text(str(data.get("ideogram_render_style", "")), 240),
|
|
"ideogram_exact_text": _sanitize_prompt_text(
|
|
str(data.get("ideogram_exact_text", "")),
|
|
"",
|
|
1000,
|
|
strip_thinking=True,
|
|
strip_markdown=True,
|
|
),
|
|
"ideogram_json_output": bool(data.get("ideogram_json_output", True)),
|
|
"negative_prompt_mode": _normalize_negative_prompt_mode(str(data.get("negative_prompt_mode", "auto"))),
|
|
"negative_prompt_guidance": _negative_prompt_guidance_text(str(data.get("negative_prompt_guidance", ""))),
|
|
}
|
|
|
|
|
|
def _make_target_profile_config(
|
|
target_profile: str = "ltx",
|
|
audio_mode: str = "auto_scene_audio",
|
|
audio_guidance: str = "",
|
|
target_duration_seconds: float = 0.0,
|
|
ltx_style: str = "",
|
|
ideogram_aspect_ratio: str = "1:1",
|
|
ideogram_render_style: str = "",
|
|
ideogram_exact_text: str = "",
|
|
ideogram_json_output: bool = True,
|
|
negative_prompt_mode: str = "auto",
|
|
negative_prompt_guidance: str = "",
|
|
) -> TargetProfileConfig:
|
|
data = _target_profile_config_to_dict(
|
|
{
|
|
"target_profile": target_profile,
|
|
"audio_mode": audio_mode,
|
|
"audio_guidance": audio_guidance,
|
|
"target_duration_seconds": target_duration_seconds,
|
|
"ltx_style": ltx_style,
|
|
"ideogram_aspect_ratio": ideogram_aspect_ratio,
|
|
"ideogram_render_style": ideogram_render_style,
|
|
"ideogram_exact_text": ideogram_exact_text,
|
|
"ideogram_json_output": ideogram_json_output,
|
|
"negative_prompt_mode": negative_prompt_mode,
|
|
"negative_prompt_guidance": negative_prompt_guidance,
|
|
}
|
|
)
|
|
return TargetProfileConfig(**data)
|
|
|
|
|
|
def _normalize_thinking_mode(thinking_mode: str) -> str:
|
|
value = (thinking_mode or "").strip().lower().replace("-", "_").replace(" ", "_")
|
|
if value in {"auto", "on", "off"}:
|
|
return value
|
|
return "auto"
|
|
|
|
|
|
def _thinking_instruction(thinking_mode: str) -> str:
|
|
mode = _normalize_thinking_mode(thinking_mode)
|
|
if mode == "off":
|
|
return "Do not include chain-of-thought or hidden reasoning; return only the final JSON."
|
|
if mode == "on":
|
|
return "You may write concise reasoning in a <think>...</think> block before the final JSON. The final JSON must still be complete and directly usable."
|
|
return "Use concise private reasoning only if it improves the prompt plan; put any reasoning in <think>...</think> before the final JSON."
|
|
|
|
|
|
def _split_reasoning_and_answer(raw_output: str) -> tuple[str, str]:
|
|
text = str(raw_output or "")
|
|
delimited_answer, has_final_boundary = _extract_delimited_final_payload(text)
|
|
reasoning_parts: list[str] = []
|
|
for match in re.finditer(r"<think>(.*?)</think>", text, flags=re.IGNORECASE | re.DOTALL):
|
|
part = _sanitize_prompt_text(match.group(1), "", 4000, strip_thinking=False, strip_markdown=True)
|
|
if part:
|
|
reasoning_parts.append(part)
|
|
for match in re.finditer(r"<\|thought\|>(.*?)(?=<\|(?:end_thought|final_json|answer|final)\|>)", text, flags=re.IGNORECASE | re.DOTALL):
|
|
part = _sanitize_prompt_text(match.group(1), "", 4000, strip_thinking=False, strip_markdown=True)
|
|
if part:
|
|
reasoning_parts.append(part)
|
|
for match in re.finditer(r"<\|channel\|>thought(.*?)(?=<\|channel\|>|$)", text, flags=re.IGNORECASE | re.DOTALL):
|
|
part = _sanitize_prompt_text(match.group(1), "", 4000, strip_thinking=False, strip_markdown=True)
|
|
if part:
|
|
reasoning_parts.append(part)
|
|
if has_final_boundary:
|
|
return (_single_paragraph(" ".join(reasoning_parts)), delimited_answer.strip())
|
|
answer = _strip_thinking(text)
|
|
first_json = answer.find("{")
|
|
if first_json > 0:
|
|
preface = _sanitize_prompt_text(answer[:first_json], "", 2000, strip_thinking=False, strip_markdown=True)
|
|
if preface:
|
|
reasoning_parts.append(preface)
|
|
answer = answer[first_json:]
|
|
return (_single_paragraph(" ".join(reasoning_parts)), answer.strip())
|
|
|
|
|
|
def _runtime_supports_pixels(config: RuntimeConfig) -> bool:
|
|
return bool(config.status.get("supports_pixels")) if isinstance(config, RuntimeConfig) else False
|
|
|
|
|
|
def _media_policy_notes(config: RuntimeConfig, context: GemmaContext | None) -> list[str]:
|
|
if not isinstance(context, GemmaContext) or context.images is None:
|
|
return []
|
|
if _runtime_supports_pixels(config):
|
|
return ["Media pixels will be passed to the active backend."]
|
|
return ["Active backend is metadata-only for image/video; sampled pixels were not sent to DiffusionGemma runtime."]
|
|
|
|
|
|
def _image_batch_to_pil(images: Any) -> list[Any]:
|
|
if images is None:
|
|
return []
|
|
try:
|
|
import numpy as np
|
|
from PIL import Image
|
|
import torch
|
|
|
|
tensor = images.detach().cpu() if isinstance(images, torch.Tensor) else images
|
|
if hasattr(tensor, "numpy"):
|
|
arr = tensor.numpy()
|
|
else:
|
|
arr = np.asarray(tensor)
|
|
if arr.ndim == 3:
|
|
arr = arr[None, ...]
|
|
pil_images = []
|
|
for frame in arr:
|
|
frame = np.clip(frame[..., :3] * 255.0, 0, 255).astype("uint8")
|
|
pil_images.append(Image.fromarray(frame))
|
|
return pil_images
|
|
except Exception:
|
|
return []
|
|
|
|
|
|
def _image_batch_hwc_shape(images: Any) -> tuple[int, int, int] | None:
|
|
shape = getattr(images, "shape", None)
|
|
if shape is None:
|
|
return None
|
|
try:
|
|
dims = tuple(int(item) for item in shape)
|
|
except Exception:
|
|
return None
|
|
if len(dims) >= 4:
|
|
return dims[-3], dims[-2], dims[-1]
|
|
if len(dims) == 3:
|
|
return dims[0], dims[1], dims[2]
|
|
return None
|
|
|
|
|
|
def _fit_reference_image_to_video_frames(reference_image: Any, video_frames: Any) -> tuple[Any, dict[str, Any]]:
|
|
report: dict[str, Any] = {"reference_image_backend_attached": False}
|
|
try:
|
|
first_reference = reference_image[:1] if len(getattr(reference_image, "shape", ())) >= 4 else reference_image
|
|
except Exception:
|
|
first_reference = reference_image
|
|
|
|
reference_shape = _image_batch_hwc_shape(first_reference)
|
|
video_shape = _image_batch_hwc_shape(video_frames)
|
|
if reference_shape:
|
|
report["reference_image_original_height"] = reference_shape[0]
|
|
report["reference_image_original_width"] = reference_shape[1]
|
|
if video_shape:
|
|
report["reference_image_backend_height"] = video_shape[0]
|
|
report["reference_image_backend_width"] = video_shape[1]
|
|
if not reference_shape or not video_shape:
|
|
return first_reference, report
|
|
if reference_shape[:2] == video_shape[:2] and reference_shape[2] == video_shape[2]:
|
|
try:
|
|
import torch
|
|
|
|
if torch.is_tensor(first_reference) and torch.is_tensor(video_frames):
|
|
aligned = first_reference
|
|
if aligned.ndim == 3 and video_frames.ndim == 4:
|
|
aligned = aligned.unsqueeze(0)
|
|
first_reference = aligned.to(device=video_frames.device, dtype=video_frames.dtype)
|
|
except Exception:
|
|
pass
|
|
report["reference_image_resized_for_video_control"] = False
|
|
return first_reference, report
|
|
|
|
try:
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
if not torch.is_tensor(first_reference) or not torch.is_tensor(video_frames):
|
|
report["reference_image_resize_error"] = "reference image or video frames are not torch tensors"
|
|
return first_reference, report
|
|
|
|
target_h, target_w, target_c = video_shape
|
|
ref = first_reference
|
|
if ref.ndim == 3:
|
|
ref = ref.unsqueeze(0)
|
|
ref = ref[:1].to(device=video_frames.device, dtype=video_frames.dtype)
|
|
if int(ref.shape[-1]) > target_c:
|
|
ref = ref[..., :target_c]
|
|
elif int(ref.shape[-1]) < target_c:
|
|
channel_pad = target_c - int(ref.shape[-1])
|
|
ref = F.pad(ref, (0, channel_pad, 0, 0, 0, 0), mode="constant", value=0.0)
|
|
|
|
nchw = ref.permute(0, 3, 1, 2).float()
|
|
src_h = max(1, int(nchw.shape[-2]))
|
|
src_w = max(1, int(nchw.shape[-1]))
|
|
scale = min(float(target_h) / float(src_h), float(target_w) / float(src_w))
|
|
resized_h = max(1, min(int(target_h), int(round(src_h * scale))))
|
|
resized_w = max(1, min(int(target_w), int(round(src_w * scale))))
|
|
resized = F.interpolate(nchw, size=(resized_h, resized_w), mode="bilinear", align_corners=False)
|
|
pad_top = max(0, (int(target_h) - resized_h) // 2)
|
|
pad_bottom = max(0, int(target_h) - resized_h - pad_top)
|
|
pad_left = max(0, (int(target_w) - resized_w) // 2)
|
|
pad_right = max(0, int(target_w) - resized_w - pad_left)
|
|
fitted = F.pad(resized, (pad_left, pad_right, pad_top, pad_bottom), mode="replicate")
|
|
fitted = fitted[..., : int(target_h), : int(target_w)]
|
|
prepared = fitted.permute(0, 2, 3, 1).to(dtype=video_frames.dtype)
|
|
report["reference_image_resized_for_video_control"] = True
|
|
report["reference_image_resize_mode"] = "aspect_fit_replicate_pad"
|
|
return prepared, report
|
|
except Exception as exc:
|
|
report["reference_image_resize_error"] = str(exc)
|
|
return first_reference, report
|
|
|
|
|
|
def _build_model_prompt(
|
|
user_prompt: str,
|
|
master_prompt: str,
|
|
target_profile: str,
|
|
media_metadata: dict[str, Any],
|
|
audio_mode: str = "auto_scene_audio",
|
|
audio_guidance: str = "",
|
|
target_duration_seconds: float = 0.0,
|
|
ltx_style: str = "",
|
|
ideogram_aspect_ratio: str = "1:1",
|
|
ideogram_render_style: str = "",
|
|
ideogram_exact_text: str = "",
|
|
ideogram_json_output: bool = True,
|
|
creativity_mode: str = "editorial",
|
|
creative_strength: float = 0.6,
|
|
negative_prompt_mode: str = "auto",
|
|
negative_prompt_guidance: str = "",
|
|
thinking_mode: str = "off",
|
|
) -> str:
|
|
target_profile = _normalize_target_profile(target_profile)
|
|
audio_mode = _normalize_audio_mode(audio_mode)
|
|
audio_guidance = _audio_guidance_text(audio_guidance)
|
|
audio_instruction = _audio_instruction(audio_mode, audio_guidance)
|
|
negative_prompt_mode = _normalize_negative_prompt_mode(negative_prompt_mode)
|
|
negative_prompt_guidance = _negative_prompt_guidance_text(negative_prompt_guidance)
|
|
negative_instruction = _negative_prompt_instruction(negative_prompt_mode, negative_prompt_guidance)
|
|
thinking_mode = _normalize_thinking_mode(thinking_mode)
|
|
thinking_instruction = _thinking_instruction(thinking_mode)
|
|
controls = _target_context(
|
|
media_metadata,
|
|
target_duration_seconds,
|
|
ltx_style,
|
|
ideogram_aspect_ratio,
|
|
ideogram_render_style,
|
|
ideogram_exact_text,
|
|
ideogram_json_output,
|
|
creativity_mode,
|
|
creative_strength,
|
|
negative_prompt_mode,
|
|
negative_prompt_guidance,
|
|
)
|
|
creativity_instruction = _creativity_instruction(creativity_mode, creative_strength)
|
|
structure_anchor = _target_structure_anchor(target_profile, ideogram_json_output)
|
|
media_grounding_mode = _media_grounding_mode(media_metadata)
|
|
media_grounding_instruction = _media_grounding_instruction(media_metadata)
|
|
ideogram_shape: Any
|
|
if ideogram_json_output:
|
|
ideogram_shape = {
|
|
"high_level_description": "One or two sentence summary of the whole image.",
|
|
"style_description": {
|
|
"aesthetics": "Aesthetic keywords.",
|
|
"lighting": "Lighting description including source time of day, ambient brightness, light sources, shadows, contrast, exposure level, blown highlights, bloom, halation, and dynamic range.",
|
|
"photo": "Use this key only for photographic outputs; include lens/camera language such as focal length impression, depth of field, bokeh, lens blur, motion blur, grain/noise, perspective compression, distortion, and camera angle; omit art_style when photo is present.",
|
|
"medium": "photograph",
|
|
"color_palette": ["#RRGGBB optional uppercase hex values"],
|
|
},
|
|
"compositional_deconstruction": {
|
|
"background": "Required background/environment description including source time of day, lighting state, exposure behavior, lens blur/bokeh, and background sharpness.",
|
|
"elements": [
|
|
{
|
|
"type": "obj",
|
|
"bbox": [0, 0, 1000, 1000],
|
|
"desc": "Object/subject description, including focus, subject sharpness, overexposed highlights, flash/direct light, shadow detail, lens blur, or motion blur when visible. Omit bbox if uncertain.",
|
|
},
|
|
{
|
|
"type": "text",
|
|
"bbox": [0, 0, 1000, 1000],
|
|
"text": "Literal text to render exactly.",
|
|
"desc": "Typography and placement description. Omit bbox if uncertain.",
|
|
}
|
|
],
|
|
},
|
|
}
|
|
else:
|
|
ideogram_shape = (
|
|
"One detailed Ideogram4 prompt sentence or paragraph with aspect ratio, style, exact readable text, "
|
|
"layout, typography, colors, and negative constraints. If the user brief is short, expand it into a complete image prompt with subject, background environment, composition, lighting, style, color palette, and text placement when relevant. Preserve source time of day, lighting state, exposure, lens blur, bokeh, depth of field, focal length impression, motion blur, grain/noise, contrast, and camera angle exactly when image/video evidence is supplied."
|
|
)
|
|
if target_profile == "ltx":
|
|
ideogram_shape = "Return an empty string for LTX target profile."
|
|
ltx_shape: Any = (
|
|
"Final LTX positive prompt only. Single continuous natural-English paragraph, usually 80-180 words. "
|
|
"Start directly with style if useful, then the action. Use active present/progressive tense, chronological flow, literal visible/audible detail, subject appearance, setting, motion, lighting, and camera behavior. "
|
|
"If image_identity_video_control synthesis is active, describe the reference-image subject performing the video action/control structure; the image controls identity and appearance, including face, hair, body type, wardrobe, accessories, styling, and distinguishing visual traits, while the video controls pose, depth, canny/edge layout, blocking, motion, timing, composition, camera choreography, and scene geometry. "
|
|
"When this synthesis mode is active, do not copy the control-video subject's clothing, hair, face, body type, accessories, or styling unless explicitly requested, and do not let the person in the control video override the identity or appearance from the reference image. "
|
|
"If the user brief is short, expand it into a complete video prompt with subject, background environment, lighting, camera framing, lens/focus, pacing, and concrete action beats. "
|
|
"When image/video evidence is supplied, preserve source time of day, ambient brightness, visible light sources, shadows, contrast, exposure, overexposed or blown highlights, flash/direct light, bloom, halation, lens blur, bokeh, depth of field, focal length impression, perspective compression, lens distortion, motion blur, grain/noise, white balance, dynamic range, and camera angle exactly. "
|
|
"When video evidence is supplied, extract camera choreography before finalizing: first-frame shot size and angle, camera height and support, subject-camera relationship, movement type and direction, zoom/dolly/truck/boom/pan/tilt/roll/orbit/rotation, focus behavior, lens/depth changes, foreground-background parallax, and final frame. Preserve the opening framing, any visible camera movement or stable hold over time, and the closing framing inside the same paragraph. If the camera performs a compound move such as zooming out while rotating 180 degrees and following the subject across a room, describe it as one continuous camera-subject relationship with start and end framing. If the camera is mostly static, say that rather than inventing a zoom, pan, reveal, exit, or wider pull-back. "
|
|
"If input is vague, add concrete cinematic details without adding unrequested characters. Mention camera movement only if requested or evident from supplied media. "
|
|
"If speech is requested, preserve exact quoted words; otherwise do not invent dialogue. If audio is enabled, weave concrete sound into the timeline beside the action. "
|
|
"No headings, timestamps, JSON, Markdown, shot labels unless requested, non-visual senses, or phrases like 'The scene opens with'."
|
|
)
|
|
if target_profile == "ideogram4":
|
|
ltx_shape = "Return an empty string for Ideogram4 target profile."
|
|
contract = {
|
|
"ltx_prompt": ltx_shape,
|
|
"ideogram_prompt": ideogram_shape,
|
|
"negative_prompt": "Final negative prompt string following the requested negative prompt mode.",
|
|
"scene_segments": [
|
|
{
|
|
"index": 0,
|
|
"duration_seconds": 0,
|
|
"prompt": "Optional final segment prompt text, not instructions.",
|
|
}
|
|
],
|
|
"metadata": {
|
|
"observations": ["Concise visual/video facts used to build prompts."],
|
|
"target_profile": target_profile,
|
|
"audio_mode": audio_mode,
|
|
"controls": controls,
|
|
},
|
|
}
|
|
return (
|
|
f"{master_prompt.strip()}\n\n"
|
|
"Return the final answer as JSON matching this shape:\n"
|
|
f"{_json_dumps(contract)}\n\n"
|
|
"Critical output rules:\n"
|
|
"- Prompt fields must be ready to paste directly into LTX-2.3 or Ideogram4.\n"
|
|
"- Interpret User raw prompt as the creative brief for the target output, even when it is only a subject, moment, style, or rough idea. Supply the missing generator-prompt details yourself.\n"
|
|
"- Only write the selected target profile. For ltx, return ideogram_prompt as an empty string. For ideogram4, return ltx_prompt as an empty string and scene_segments as an empty list.\n"
|
|
"- Do not put writing instructions, schema explanations, or adapter instructions inside any prompt field.\n"
|
|
f"- Final response boundary: start the final answer with {FINAL_JSON_OPEN}, write one valid JSON object, close with {FINAL_JSON_CLOSE}, and write nothing after the closing marker.\n"
|
|
"- Use duration as pacing guidance for the LTX prompt, but do not write 'Target duration' or duration metadata in the prompt text.\n"
|
|
"- Ideogram JSON must follow its caption schema: top-level high_level_description, style_description, compositional_deconstruction; compositional_deconstruction contains background then elements; elements use type=obj/text and desc/text keys; no unknown layout/aspect_ratio keys inside the caption.\n"
|
|
"- Ideogram color_palette values must be uppercase #RRGGBB hex strings if included; omit color_palette when no hex palette is known.\n"
|
|
"- Use Ideogram aspect ratio as layout guidance only; do not include aspect_ratio as a JSON caption key.\n"
|
|
"- Preserve exact quoted or user-specified text verbatim for Ideogram text elements.\n"
|
|
"- Apply creative direction by adding specific visual details, not by adding meta commentary about creativity.\n"
|
|
"- Respect negative prompt mode exactly; custom negatives must be returned unchanged, and empty mode must return an empty string.\n"
|
|
"- Do not invent dialogue unless the user asked for speech, talking, singing, or conversation.\n\n"
|
|
"- Preserve source lighting and time of day as non-negotiable visual identity when media or visual_description is supplied. If the source is night, the final prompt must repeatedly and concretely say nighttime/after dark, visible artificial light sources, deep shadows, and no daylight/sunlit/blue-sky conditions unless those are visibly present.\n"
|
|
"- Preserve source camera/capture language as non-negotiable visual identity when media or visual_description is supplied. Use concrete industry terms when visible or strongly implied: overexposed highlights, clipped whites, underexposed shadows, direct flash, bloom, halation, shallow depth of field, soft focus, lens blur, bokeh, motion blur, focal length impression, wide-angle distortion, telephoto compression, high ISO grain/noise, white balance, high contrast, low dynamic range, and camera angle.\n"
|
|
"- If media_synthesis_mode is image_identity_video_control, the reference image controls subject identity and appearance: face, hair, body type, wardrobe, accessories, styling, and distinguishing visual traits. The video controls pose, depth, canny/edge layout, composition, motion, timing, camera choreography, blocking, and scene geometry. The LTX prompt must describe the reference-image subject performing the video action/control structure.\n"
|
|
"- In image_identity_video_control, do not copy the control-video subject's wardrobe, hair, face, body type, accessories, or styling into the LTX prompt unless the user explicitly asks for those video-subject traits. If the image subject and video subject differ, replace the video subject appearance with the reference-image appearance.\n"
|
|
"- For LTX with source video evidence, perform camera-choreography extraction before finalizing: first-frame framing, shot size, camera height/angle, camera support or stability, subject-camera relationship, movement type and direction, zoom/dolly/truck/boom/pan/tilt/roll/orbit/rotation, focus behavior, lens/depth changes, parallax, and final frame.\n"
|
|
"- Preserve complex visible camera moves as continuous relationships, not generic labels. For example, if the source zooms out while rotating 180 degrees and tracking the actress to the other side of the room, the LTX prompt should say that plainly with where the camera starts, how it moves relative to her, and where it settles.\n"
|
|
"- For LTX with source video evidence, preserve the clip chronologically by keeping the opening framing, any visible camera motion or stable hold, and the closing framing. If the source does not visibly zoom, pan, reveal, pull back, widen, rotate, orbit, or stage a new action beat, do not invent one.\n"
|
|
"- If Media grounding mode is metadata_only, do not claim to see or describe the supplied image/video pixels. Dimensions, frame counts, and sample indices are not visual evidence.\n"
|
|
"- If Media grounding mode is visual_description, treat the visual_description as the only visual evidence and make the final prompt specific to it.\n"
|
|
"- If Media grounding mode is pixels, ground observations in the attached pixels.\n\n"
|
|
f"Target profile: {target_profile}\n"
|
|
f"Target controls: {_json_dumps(controls)}\n"
|
|
f"Media grounding mode: {media_grounding_mode}\n"
|
|
f"Media grounding instruction: {media_grounding_instruction}\n"
|
|
f"Creative direction: {creativity_instruction}\n"
|
|
f"Thinking mode: {thinking_mode}\n"
|
|
f"Thinking instruction: {thinking_instruction}\n"
|
|
f"Audio mode: {audio_mode}\n"
|
|
f"Audio instruction: {audio_instruction}\n"
|
|
f"Audio guidance: {audio_guidance or 'none'}\n"
|
|
f"Negative prompt mode: {negative_prompt_mode}\n"
|
|
f"Negative prompt instruction: {negative_instruction}\n"
|
|
f"Negative prompt guidance: {negative_prompt_guidance or 'none'}\n"
|
|
f"Media metadata: {_json_dumps(media_metadata)}\n"
|
|
f"User raw prompt: {user_prompt.strip()}\n"
|
|
f"Final structure anchor: {structure_anchor}\n"
|
|
f"Final response boundary: {FINAL_JSON_OPEN} JSON_OBJECT {FINAL_JSON_CLOSE}"
|
|
)
|
|
|
|
|
|
def _template_packet(
|
|
user_prompt: str,
|
|
target_profile: str,
|
|
media_metadata: dict[str, Any],
|
|
duration_seconds: float,
|
|
audio_mode: str = "auto_scene_audio",
|
|
audio_guidance: str = "",
|
|
target_duration_seconds: float = 0.0,
|
|
ltx_style: str = "",
|
|
ideogram_aspect_ratio: str = "1:1",
|
|
ideogram_render_style: str = "",
|
|
ideogram_exact_text: str = "",
|
|
ideogram_json_output: bool = True,
|
|
creativity_mode: str = "editorial",
|
|
creative_strength: float = 0.6,
|
|
negative_prompt_mode: str = "auto",
|
|
negative_prompt_guidance: str = "",
|
|
) -> dict[str, Any]:
|
|
target_profile = _normalize_target_profile(target_profile)
|
|
negative_prompt_mode = _normalize_negative_prompt_mode(negative_prompt_mode)
|
|
negative_prompt_guidance = _negative_prompt_guidance_text(negative_prompt_guidance)
|
|
prompt_base = _clean_prompt_request_text(user_prompt)
|
|
visual_description = _visual_description_from_metadata(media_metadata)
|
|
media_grounding_mode = _media_grounding_mode(media_metadata)
|
|
limitation_prompt = False
|
|
if visual_description:
|
|
if prompt_base and not _is_media_reference_only_request(user_prompt):
|
|
base = _single_paragraph(f"{visual_description}. User direction: {prompt_base}")
|
|
else:
|
|
base = visual_description
|
|
elif media_grounding_mode == "pixels" and _is_media_reference_only_request(user_prompt):
|
|
limitation_prompt = True
|
|
base = (
|
|
"DiffusionGemma received media pixels, but it did not return usable grounded prompt text. "
|
|
"Re-run with runtime_required enabled to surface backend errors, or add a visual_description caption node."
|
|
)
|
|
elif media_grounding_mode == "metadata_only" and _is_media_reference_only_request(user_prompt):
|
|
limitation_prompt = True
|
|
base = (
|
|
"Media pixels were supplied, but the active DiffusionGemma backend cannot inspect them. "
|
|
"Add a visual_description from an image or video caption node, or use a pixel-capable backend."
|
|
)
|
|
else:
|
|
base = prompt_base or "A clear cinematic scene with a distinct subject, setting, style, and motion."
|
|
controls = _target_context(
|
|
media_metadata,
|
|
target_duration_seconds,
|
|
ltx_style,
|
|
ideogram_aspect_ratio,
|
|
ideogram_render_style,
|
|
ideogram_exact_text,
|
|
ideogram_json_output,
|
|
creativity_mode,
|
|
creative_strength,
|
|
negative_prompt_mode,
|
|
negative_prompt_guidance,
|
|
)
|
|
art_direction = _creative_art_direction_phrase(creativity_mode, creative_strength)
|
|
effective_duration = _effective_duration_seconds(media_metadata, target_duration_seconds) or duration_seconds
|
|
media_hint = ""
|
|
if media_grounding_mode == "visual_description":
|
|
media_hint = "The prompt is grounded in the supplied visual description rather than direct pixel analysis."
|
|
elif media_grounding_mode == "pixels" and _is_image_identity_video_control(media_metadata):
|
|
count = media_metadata.get("video_sampled_frame_count", media_metadata.get("sampled_frame_count", 0))
|
|
media_hint = (
|
|
f"The output subject identity comes from the reference image, while the action/control structure stays continuous across {count} sampled video frames. "
|
|
"Use the reference image for face, hair, body type, wardrobe, accessories, styling, and distinguishing visual traits. "
|
|
"Preserve the video pose, depth/canny-style layout, blocking, camera behavior, opening framing, visible movement or stable hold, closing framing, source exposure, lens/depth-of-field character, blur, grain, contrast, and lighting state. "
|
|
"Do not copy the control-video subject's wardrobe, hair, face, body type, accessories, or styling, and do not let the person in the control video override the reference-image identity or appearance."
|
|
)
|
|
elif media_grounding_mode == "pixels" and media_metadata.get("source") == "video":
|
|
count = media_metadata.get("sampled_frame_count", 0)
|
|
media_hint = (
|
|
f"The action stays continuous across the sampled video evidence from {count} frames, preserving subject motion, camera behavior, "
|
|
"opening framing, any visible camera movement or stable hold, closing framing, source exposure, lens/depth-of-field character, blur, grain, contrast, and lighting state. "
|
|
"Do not invent zooms, pans, reveals, exits, or wider pull-backs unless they are visibly present."
|
|
)
|
|
elif media_grounding_mode == "pixels" and media_metadata.get("source") in {"image", "image+video"}:
|
|
media_hint = (
|
|
"The composition preserves the supplied image evidence, visual identity, layout, style, exposure, lens/depth-of-field character, "
|
|
"blur, grain, contrast, and lighting state."
|
|
)
|
|
|
|
audio_mode = _normalize_audio_mode(audio_mode)
|
|
audio_guidance = _audio_guidance_text(audio_guidance)
|
|
audio_hint = ""
|
|
if audio_guidance and audio_mode != "visual_only":
|
|
audio_hint = f"Scene-specific sound design features {audio_guidance}."
|
|
elif audio_mode == "explicit_sound_design":
|
|
audio_hint = "Concrete diegetic sound follows the visible action and setting."
|
|
|
|
ltx_clauses: list[str] = []
|
|
if controls["ltx_style"] and not limitation_prompt:
|
|
ltx_clauses.append(f"Style: {controls['ltx_style']}.")
|
|
ltx_clauses.append(_ensure_terminal_period(base))
|
|
if media_hint and not limitation_prompt:
|
|
ltx_clauses.append(media_hint)
|
|
if art_direction and not limitation_prompt:
|
|
ltx_clauses.append(art_direction)
|
|
if audio_hint and not limitation_prompt:
|
|
ltx_clauses.append(audio_hint)
|
|
ltx = _single_paragraph(" ".join(ltx_clauses))
|
|
synthesis_metadata: dict[str, Any] = {}
|
|
_copy_synthesis_metadata(synthesis_metadata, media_metadata)
|
|
|
|
if ideogram_json_output:
|
|
ideogram: Any = _ideogram_prompt_object(
|
|
base,
|
|
ideogram_aspect_ratio,
|
|
ideogram_render_style,
|
|
controls["ideogram_exact_text"],
|
|
creativity_mode,
|
|
creative_strength,
|
|
)
|
|
else:
|
|
ideogram = _ideogram_prompt_text(
|
|
base,
|
|
ideogram_aspect_ratio,
|
|
ideogram_render_style,
|
|
controls["ideogram_exact_text"],
|
|
)
|
|
if target_profile == "ltx":
|
|
ideogram = ""
|
|
elif target_profile == "ideogram4":
|
|
ltx = ""
|
|
segments = _split_segments(ltx, effective_duration) if ltx else []
|
|
return {
|
|
"ltx_prompt": ltx,
|
|
"ideogram_prompt": ideogram,
|
|
"negative_prompt": _apply_negative_prompt_policy("", negative_prompt_mode, negative_prompt_guidance),
|
|
"scene_segments": segments,
|
|
"metadata": {
|
|
"backend": "template",
|
|
"fallback_reason": "No compatible DiffusionGemma runtime was used.",
|
|
"target_profile": target_profile,
|
|
"visual_grounding_mode": media_grounding_mode,
|
|
"media": media_metadata,
|
|
"audio_mode": audio_mode,
|
|
"audio_guidance": audio_guidance,
|
|
"negative_prompt_mode": negative_prompt_mode,
|
|
"negative_prompt_guidance": negative_prompt_guidance,
|
|
"controls": controls,
|
|
"ready_for_generation": False,
|
|
"blocked_reasons": ["template_fallback"],
|
|
**synthesis_metadata,
|
|
},
|
|
}
|
|
|
|
|
|
def _extract_labeled_prompt(text: str, label: str) -> str:
|
|
labels = r"(?:ltx_prompt|ideogram_prompt|negative_prompt|scene_segments|metadata)"
|
|
pattern = rf"['\"]?\b{re.escape(label)}\b['\"]?\s*[:=]\s*(.*?)(?=['\"]?\b{labels}\b['\"]?\s*[:=]|$)"
|
|
match = re.search(pattern, text, flags=re.IGNORECASE | re.DOTALL)
|
|
if not match:
|
|
return ""
|
|
value = match.group(1).strip().rstrip(",")
|
|
if value.startswith('"') and value.endswith('"'):
|
|
try:
|
|
decoded = json.loads(value)
|
|
if isinstance(decoded, str):
|
|
value = decoded
|
|
except Exception:
|
|
value = value.strip('"')
|
|
return _sanitize_prompt_text(value, "", 8000)
|
|
|
|
|
|
def _has_salvageable_plain_text(text: str) -> bool:
|
|
cleaned = _strip_markdown(_strip_thinking(text or "")).strip()
|
|
if not cleaned:
|
|
return False
|
|
if cleaned in {"{", "}", "[", "]", "{}", "[]"}:
|
|
return False
|
|
letters = re.findall(r"[A-Za-z]", cleaned)
|
|
words = re.findall(r"[A-Za-z]{2,}", cleaned)
|
|
return len(letters) >= 20 and len(words) >= 4
|
|
|
|
|
|
def _packet_from_plain_text_output(
|
|
raw_output: str,
|
|
user_prompt: str,
|
|
target_profile: str,
|
|
media_metadata: dict[str, Any],
|
|
duration_seconds: float,
|
|
audio_mode: str,
|
|
audio_guidance: str,
|
|
target_duration_seconds: float,
|
|
ltx_style: str,
|
|
ideogram_aspect_ratio: str,
|
|
ideogram_render_style: str,
|
|
ideogram_exact_text: str,
|
|
ideogram_json_output: bool,
|
|
creativity_mode: str,
|
|
creative_strength: float,
|
|
negative_prompt_mode: str,
|
|
negative_prompt_guidance: str,
|
|
max_output_chars: int,
|
|
) -> dict[str, Any] | None:
|
|
text = _sanitize_prompt_text(raw_output, "", int(max_output_chars))
|
|
if not text:
|
|
return None
|
|
|
|
ltx_labeled = _extract_labeled_prompt(text, "ltx_prompt")
|
|
ideogram_labeled = _extract_labeled_prompt(text, "ideogram_prompt")
|
|
negative_labeled = _extract_labeled_prompt(text, "negative_prompt")
|
|
if not (ltx_labeled or ideogram_labeled or _has_salvageable_plain_text(text)):
|
|
return None
|
|
template = _template_packet(
|
|
user_prompt,
|
|
target_profile,
|
|
media_metadata,
|
|
duration_seconds,
|
|
audio_mode,
|
|
audio_guidance,
|
|
target_duration_seconds,
|
|
ltx_style,
|
|
ideogram_aspect_ratio,
|
|
ideogram_render_style,
|
|
ideogram_exact_text,
|
|
ideogram_json_output,
|
|
creativity_mode,
|
|
creative_strength,
|
|
negative_prompt_mode,
|
|
negative_prompt_guidance,
|
|
)
|
|
profile = _normalize_target_profile(target_profile)
|
|
plain = _clean_prompt_request_text(ltx_labeled or ideogram_labeled or text)
|
|
if profile == "ideogram4":
|
|
template["ideogram_prompt"] = ideogram_labeled or plain
|
|
else:
|
|
template["ltx_prompt"] = ltx_labeled or plain
|
|
if ideogram_labeled:
|
|
template["ideogram_prompt"] = ideogram_labeled
|
|
template["negative_prompt"] = _apply_negative_prompt_policy(
|
|
negative_labeled or template.get("negative_prompt", ""),
|
|
negative_prompt_mode,
|
|
negative_prompt_guidance,
|
|
)
|
|
template["scene_segments"] = _split_segments(
|
|
str(template.get("ltx_prompt", "")),
|
|
_effective_duration_seconds(media_metadata, target_duration_seconds) or duration_seconds,
|
|
)
|
|
metadata = template.setdefault("metadata", {})
|
|
metadata.pop("fallback_reason", None)
|
|
metadata.update(
|
|
{
|
|
"plain_text_salvage": True,
|
|
"json_parse_warning": "Model output did not contain valid JSON; plain text output was used.",
|
|
"template_used_for_missing_fields": True,
|
|
}
|
|
)
|
|
return template
|
|
|
|
|
|
def repair_or_salvage_prompt_packet(
|
|
raw_output: str,
|
|
fallback_context: GemmaContext | str | None,
|
|
target_config: TargetProfileConfig | dict[str, Any] | None,
|
|
max_output_chars: int = 8000,
|
|
creativity_mode: str = "editorial",
|
|
creative_strength: float = 0.6,
|
|
) -> tuple[dict[str, Any], bool, str]:
|
|
context = fallback_context if isinstance(fallback_context, GemmaContext) else _gemma_context_from_media(str(fallback_context or ""))
|
|
target = _target_profile_config_to_dict(target_config)
|
|
media_metadata = context.media_metadata.copy()
|
|
duration = _safe_float(media_metadata.get("duration_seconds"), 0.0)
|
|
parsed = _extract_json_object(raw_output)
|
|
if parsed is not None:
|
|
return parsed, True, ""
|
|
salvaged = _packet_from_plain_text_output(
|
|
raw_output,
|
|
context.user_prompt,
|
|
target["target_profile"],
|
|
media_metadata,
|
|
duration,
|
|
target["audio_mode"],
|
|
target["audio_guidance"],
|
|
target["target_duration_seconds"],
|
|
target["ltx_style"],
|
|
target["ideogram_aspect_ratio"],
|
|
target["ideogram_render_style"],
|
|
target["ideogram_exact_text"],
|
|
target["ideogram_json_output"],
|
|
creativity_mode,
|
|
creative_strength,
|
|
target["negative_prompt_mode"],
|
|
target["negative_prompt_guidance"],
|
|
int(max_output_chars),
|
|
)
|
|
if salvaged is not None:
|
|
return salvaged, False, "Model output was not valid JSON; labeled/plain text salvage was used."
|
|
fallback = _template_packet(
|
|
context.user_prompt,
|
|
target["target_profile"],
|
|
media_metadata,
|
|
duration,
|
|
target["audio_mode"],
|
|
target["audio_guidance"],
|
|
target["target_duration_seconds"],
|
|
target["ltx_style"],
|
|
target["ideogram_aspect_ratio"],
|
|
target["ideogram_render_style"],
|
|
target["ideogram_exact_text"],
|
|
target["ideogram_json_output"],
|
|
creativity_mode,
|
|
creative_strength,
|
|
target["negative_prompt_mode"],
|
|
target["negative_prompt_guidance"],
|
|
)
|
|
return fallback, False, "Model output was empty or unusable; template fallback packet was used."
|
|
|
|
|
|
def _packet_to_prompt_outputs(
|
|
packet: dict[str, Any],
|
|
context: GemmaContext,
|
|
target_config: TargetProfileConfig | dict[str, Any],
|
|
max_output_chars: int,
|
|
creativity_mode: str = "editorial",
|
|
creative_strength: float = 0.6,
|
|
) -> tuple[str, str, str, str, str, str, str, str]:
|
|
target = _target_profile_config_to_dict(target_config)
|
|
profile = target["target_profile"]
|
|
prompt_fallback = _clean_prompt_request_text(context.user_prompt)
|
|
packet_metadata = packet.get("metadata") if isinstance(packet.get("metadata"), dict) else {}
|
|
packet["metadata"] = packet_metadata
|
|
packet_media = packet_metadata.get("media") if isinstance(packet_metadata.get("media"), dict) else {}
|
|
media_metadata = context.media_metadata.copy()
|
|
if packet_media:
|
|
packet_metadata["model_reported_media"] = packet_media.copy()
|
|
if (
|
|
isinstance(context, GemmaContext)
|
|
and context.images is not None
|
|
and any(key in packet_metadata for key in ("media_policy", "runtime_ready", "raw_output_available"))
|
|
):
|
|
for key in (
|
|
"pixel_tensor_present",
|
|
"pixels_sent_to_backend",
|
|
"visual_grounding_mode",
|
|
"transformers_video_transport",
|
|
):
|
|
if key in packet_media:
|
|
media_metadata[key] = packet_media[key]
|
|
if isinstance(packet_media.get("warnings"), list):
|
|
existing = list(media_metadata.get("warnings", [])) if isinstance(media_metadata.get("warnings"), list) else []
|
|
for warning in packet_media["warnings"]:
|
|
if warning not in existing:
|
|
existing.append(warning)
|
|
media_metadata["warnings"] = existing
|
|
media_metadata["visual_grounding_mode"] = _media_grounding_mode(media_metadata)
|
|
packet_metadata["media"] = media_metadata.copy()
|
|
_copy_synthesis_metadata(packet_metadata, media_metadata)
|
|
duration = _safe_float(media_metadata.get("duration_seconds"), 0.0)
|
|
effective_duration = _effective_duration_seconds(media_metadata, target["target_duration_seconds"])
|
|
ltx_prompt = ""
|
|
if profile == "ltx":
|
|
ltx_prompt = _prompt_field_to_text(packet.get("ltx_prompt", ""), prompt_fallback, int(max_output_chars))
|
|
if _media_grounding_mode(media_metadata) == "metadata_only" and _looks_like_unseen_media_bluff(ltx_prompt):
|
|
replacement = _template_packet(
|
|
context.user_prompt,
|
|
"ltx",
|
|
media_metadata,
|
|
duration,
|
|
target["audio_mode"],
|
|
target["audio_guidance"],
|
|
target["target_duration_seconds"],
|
|
target["ltx_style"],
|
|
target["ideogram_aspect_ratio"],
|
|
target["ideogram_render_style"],
|
|
target["ideogram_exact_text"],
|
|
target["ideogram_json_output"],
|
|
creativity_mode,
|
|
creative_strength,
|
|
target["negative_prompt_mode"],
|
|
target["negative_prompt_guidance"],
|
|
)
|
|
ltx_prompt = _prompt_field_to_text(replacement.get("ltx_prompt", ""), prompt_fallback, int(max_output_chars))
|
|
packet.setdefault("metadata", {})["visual_grounding_warning"] = (
|
|
"Generic media-description output was replaced because the active backend cannot inspect pixels."
|
|
)
|
|
ideogram_value = packet.get("ideogram_prompt", "")
|
|
if not ideogram_value and {"high_level_description", "compositional_deconstruction"}.issubset(packet.keys()):
|
|
ideogram_value = packet
|
|
ideogram_prompt = ""
|
|
if profile == "ideogram4":
|
|
ideogram_prompt = _ideogram_prompt_value_to_text(
|
|
ideogram_value,
|
|
prompt_fallback,
|
|
target["ideogram_aspect_ratio"],
|
|
target["ideogram_render_style"],
|
|
target["ideogram_exact_text"],
|
|
target["ideogram_json_output"],
|
|
int(max_output_chars),
|
|
creativity_mode,
|
|
creative_strength,
|
|
)
|
|
if _media_grounding_mode(media_metadata) == "metadata_only" and _looks_like_unseen_media_bluff(ideogram_prompt):
|
|
replacement = _template_packet(
|
|
context.user_prompt,
|
|
"ideogram4",
|
|
media_metadata,
|
|
duration,
|
|
target["audio_mode"],
|
|
target["audio_guidance"],
|
|
target["target_duration_seconds"],
|
|
target["ltx_style"],
|
|
target["ideogram_aspect_ratio"],
|
|
target["ideogram_render_style"],
|
|
target["ideogram_exact_text"],
|
|
target["ideogram_json_output"],
|
|
creativity_mode,
|
|
creative_strength,
|
|
target["negative_prompt_mode"],
|
|
target["negative_prompt_guidance"],
|
|
)
|
|
ideogram_prompt = _ideogram_prompt_value_to_text(
|
|
replacement.get("ideogram_prompt", ""),
|
|
prompt_fallback,
|
|
target["ideogram_aspect_ratio"],
|
|
target["ideogram_render_style"],
|
|
target["ideogram_exact_text"],
|
|
target["ideogram_json_output"],
|
|
int(max_output_chars),
|
|
creativity_mode,
|
|
creative_strength,
|
|
)
|
|
packet.setdefault("metadata", {})["visual_grounding_warning"] = (
|
|
"Generic media-description output was replaced because the active backend cannot inspect pixels."
|
|
)
|
|
negative_prompt = _apply_negative_prompt_policy(
|
|
packet.get("negative_prompt", ""),
|
|
target["negative_prompt_mode"],
|
|
target["negative_prompt_guidance"],
|
|
)
|
|
scene_segments = packet.get("scene_segments")
|
|
if profile == "ideogram4":
|
|
scene_segments = []
|
|
elif not isinstance(scene_segments, list) or not scene_segments:
|
|
scene_segments = _split_segments(ltx_prompt, effective_duration or duration)
|
|
local_prompts, segment_lengths = _segments_to_director_strings(scene_segments)
|
|
prompt_to_verify = ideogram_prompt if profile == "ideogram4" else ltx_prompt
|
|
ready_for_generation, blocked_reasons = _packet_generation_readiness(
|
|
packet,
|
|
context,
|
|
target,
|
|
prompt_to_verify,
|
|
media_metadata,
|
|
)
|
|
packet_metadata["ready_for_generation"] = bool(ready_for_generation)
|
|
packet_metadata["blocked_reasons"] = blocked_reasons
|
|
packet_metadata["claim_verification_passed"] = not any(
|
|
reason.startswith("metadata_only_") or reason.startswith("visual_description_contradiction")
|
|
for reason in blocked_reasons
|
|
)
|
|
return (
|
|
ltx_prompt,
|
|
ideogram_prompt,
|
|
negative_prompt,
|
|
_json_dumps(scene_segments),
|
|
local_prompts,
|
|
segment_lengths,
|
|
target["ideogram_aspect_ratio"],
|
|
_json_dumps(packet.get("metadata") if isinstance(packet.get("metadata"), dict) else {}),
|
|
)
|
|
|
|
|
|
class _ComfyModelOptNVFP4Experts:
|
|
@staticmethod
|
|
def build(text_config: Any, ops: Any, device: Any):
|
|
import torch
|
|
|
|
class ComfyModelOptNVFP4Experts(torch.nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.num_experts = int(text_config.num_experts)
|
|
self.hidden_dim = int(text_config.hidden_size)
|
|
self.intermediate_dim = int(text_config.moe_intermediate_size)
|
|
self.gate_proj = ops.MoEExperts(self.num_experts, self.hidden_dim, self.intermediate_dim, bias=False, device=device)
|
|
self.up_proj = ops.MoEExperts(self.num_experts, self.hidden_dim, self.intermediate_dim, bias=False, device=device)
|
|
self.down_proj = ops.MoEExperts(self.num_experts, self.intermediate_dim, self.hidden_dim, bias=False, device=device)
|
|
from transformers.activations import ACT2FN
|
|
|
|
self.act_fn = ACT2FN[text_config.hidden_activation]
|
|
|
|
@staticmethod
|
|
def _expert_linear(bank: Any, input_tensor: Any, expert_idx: int) -> Any:
|
|
from comfy.quant_ops import QuantizedTensor
|
|
|
|
weight = getattr(bank, "weight", None)
|
|
if isinstance(weight, QuantizedTensor):
|
|
expert_weight = bank._expert_qt_from(weight, expert_idx)
|
|
input_scale = getattr(bank, "input_scale", None)
|
|
scale = None
|
|
if input_scale is not None:
|
|
scale = input_scale[expert_idx] if getattr(input_scale, "dim", lambda: 0)() else input_scale
|
|
scale = scale.to(device=input_tensor.device)
|
|
quantized_input = QuantizedTensor.from_float(input_tensor, bank.layout_type, scale=scale)
|
|
return torch.nn.functional.linear(quantized_input, expert_weight, None)
|
|
return bank.expert_linear(input_tensor, expert_idx)
|
|
|
|
def forward(self, hidden_states: Any, top_k_index: Any, top_k_weights: Any) -> Any:
|
|
final_hidden_states = torch.zeros_like(hidden_states)
|
|
with torch.no_grad():
|
|
expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
|
|
expert_mask = expert_mask.permute(2, 1, 0)
|
|
expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
|
|
|
|
for expert_idx_tensor in expert_hit:
|
|
expert_idx = int(expert_idx_tensor[0].item())
|
|
if expert_idx == self.num_experts:
|
|
continue
|
|
top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
|
|
current_state = hidden_states[token_idx]
|
|
gate = self._expert_linear(self.gate_proj, current_state, expert_idx)
|
|
up = self._expert_linear(self.up_proj, current_state, expert_idx)
|
|
current_hidden_states = self.act_fn(gate) * up
|
|
current_hidden_states = self._expert_linear(self.down_proj, current_hidden_states, expert_idx)
|
|
current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
|
|
final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
|
|
|
|
return final_hidden_states
|
|
|
|
return ComfyModelOptNVFP4Experts()
|
|
|
|
|
|
def _resolve_model_leaf(model: Any, name: str) -> tuple[Any, str]:
|
|
parts = name.split(".")
|
|
parent = model
|
|
for part in parts[:-1]:
|
|
parent = getattr(parent, part)
|
|
return parent, parts[-1]
|
|
|
|
|
|
def _model_has_leaf(model: Any, name: str) -> bool:
|
|
try:
|
|
parent, leaf = _resolve_model_leaf(model, name)
|
|
except Exception:
|
|
return False
|
|
return leaf in getattr(parent, "_parameters", {}) or leaf in getattr(parent, "_buffers", {}) or hasattr(parent, leaf)
|
|
|
|
|
|
def _set_model_leaf(model: Any, name: str, value: Any) -> None:
|
|
import torch
|
|
|
|
parent, leaf = _resolve_model_leaf(model, name)
|
|
if isinstance(value, torch.nn.Parameter):
|
|
parent._parameters[leaf] = value
|
|
elif leaf in getattr(parent, "_parameters", {}):
|
|
parent._parameters[leaf] = torch.nn.Parameter(value, requires_grad=False)
|
|
elif leaf in getattr(parent, "_buffers", {}):
|
|
parent._buffers[leaf] = value
|
|
else:
|
|
setattr(parent, leaf, value)
|
|
|
|
|
|
def _decoder_to_encoder_language_key(key: str) -> str:
|
|
if key == "model.decoder.embed_tokens.weight":
|
|
return "model.encoder.language_model.embed_tokens.weight"
|
|
if key == "model.decoder.norm.weight":
|
|
return "model.encoder.language_model.norm.weight"
|
|
if key.startswith("model.decoder.layers."):
|
|
remainder = key[len("model.decoder.layers.") :]
|
|
if ".experts." in remainder or remainder.endswith(".layer_scalar") or remainder == "layer_scalar":
|
|
return ""
|
|
return f"model.encoder.language_model.layers.{remainder}"
|
|
return ""
|
|
|
|
|
|
def _is_modelopt_expert_tensor_key(key: str) -> bool:
|
|
return bool(re.search(r"^model\.decoder\.layers\.\d+\.experts\.\d+\.(?:gate_proj|up_proj|down_proj)\.", key))
|
|
|
|
|
|
def _load_tensor_from_weight_map(model_path: str, weight_map: dict[str, str], key: str, handles: dict[str, Any]) -> Any:
|
|
from safetensors import safe_open
|
|
|
|
shard = weight_map[key]
|
|
handle = handles.get(shard)
|
|
if handle is None:
|
|
handle = safe_open(str(Path(model_path) / shard), framework="pt", device="cpu")
|
|
handles[shard] = handle
|
|
return handle.get_tensor(key)
|
|
|
|
|
|
def _install_nvfp4_bank(
|
|
bank: Any,
|
|
model_path: str,
|
|
weight_map: dict[str, str],
|
|
layer_idx: int,
|
|
projection: str,
|
|
in_features: int,
|
|
out_features: int,
|
|
device: Any,
|
|
) -> None:
|
|
import torch
|
|
from comfy.quant_ops import QuantizedTensor, get_layout_class
|
|
|
|
handles: dict[str, Any] = {}
|
|
|
|
def tensor_for(expert_idx: int, suffix: str) -> Any:
|
|
key = f"model.decoder.layers.{layer_idx}.experts.{expert_idx}.{projection}.{suffix}"
|
|
return _load_tensor_from_weight_map(model_path, weight_map, key, handles)
|
|
|
|
try:
|
|
qdata = _modelopt_to_comfy_fp4_packing(
|
|
torch.stack([tensor_for(i, "weight") for i in range(bank.num_experts)], dim=0)
|
|
).to(device=device)
|
|
block_scale = torch.stack(
|
|
[_modelopt_to_comfy_block_scale(tensor_for(i, "weight_scale")) for i in range(bank.num_experts)],
|
|
dim=0,
|
|
).to(device=device)
|
|
tensor_scale = torch.stack([tensor_for(i, "weight_scale_2") for i in range(bank.num_experts)], dim=0).to(
|
|
device=device
|
|
)
|
|
input_scale = torch.stack([tensor_for(i, "input_scale") for i in range(bank.num_experts)], dim=0).to(device=device)
|
|
finally:
|
|
handles.clear()
|
|
|
|
layout_cls = get_layout_class("TensorCoreNVFP4Layout")
|
|
params = layout_cls.Params(
|
|
scale=tensor_scale,
|
|
block_scale=block_scale,
|
|
orig_dtype=torch.bfloat16,
|
|
orig_shape=(bank.num_experts, int(out_features), int(in_features)),
|
|
)
|
|
bank.quant_format = "nvfp4"
|
|
bank.layout_type = "TensorCoreNVFP4Layout"
|
|
bank.weight = torch.nn.Parameter(
|
|
QuantizedTensor(qdata.to(torch.uint8), bank.layout_type, params),
|
|
requires_grad=False,
|
|
)
|
|
bank.input_scale = torch.nn.Parameter(input_scale, requires_grad=False)
|
|
|
|
|
|
def _replace_and_load_nvfp4_experts(model: Any, config: Any, model_path: str, weight_map: dict[str, str], device: Any) -> None:
|
|
import torch
|
|
from comfy.ops import mixed_precision_ops
|
|
|
|
ops = mixed_precision_ops({"format": "nvfp4"}, compute_dtype=torch.bfloat16)
|
|
text_config = config.text_config
|
|
total_layers = int(text_config.num_hidden_layers)
|
|
_dg_log("Loading NVFP4 expert banks: %s layers.", total_layers)
|
|
layer_iter = range(total_layers)
|
|
tqdm_bar = None
|
|
if _progress_enabled():
|
|
try:
|
|
from tqdm.auto import tqdm
|
|
|
|
tqdm_bar = tqdm(layer_iter, total=total_layers, desc="DiffusionGemma NVFP4 experts", unit="layer", leave=True)
|
|
layer_iter = tqdm_bar
|
|
except Exception:
|
|
tqdm_bar = None
|
|
for layer_idx in layer_iter:
|
|
experts = _ComfyModelOptNVFP4Experts.build(text_config, ops, device)
|
|
_install_nvfp4_bank(
|
|
experts.gate_proj,
|
|
model_path,
|
|
weight_map,
|
|
layer_idx,
|
|
"gate_proj",
|
|
int(text_config.hidden_size),
|
|
int(text_config.moe_intermediate_size),
|
|
device,
|
|
)
|
|
_install_nvfp4_bank(
|
|
experts.up_proj,
|
|
model_path,
|
|
weight_map,
|
|
layer_idx,
|
|
"up_proj",
|
|
int(text_config.hidden_size),
|
|
int(text_config.moe_intermediate_size),
|
|
device,
|
|
)
|
|
_install_nvfp4_bank(
|
|
experts.down_proj,
|
|
model_path,
|
|
weight_map,
|
|
layer_idx,
|
|
"down_proj",
|
|
int(text_config.moe_intermediate_size),
|
|
int(text_config.hidden_size),
|
|
device,
|
|
)
|
|
model.model.decoder.layers[layer_idx].experts = experts
|
|
model.model.encoder.language_model.layers[layer_idx].experts = experts
|
|
if tqdm_bar is not None:
|
|
tqdm_bar.close()
|
|
_dg_log("Loaded NVFP4 expert banks.")
|
|
|
|
|
|
def _load_modelopt_nvfp4_bridge(config: RuntimeConfig) -> tuple[Any, Any]:
|
|
import torch
|
|
from accelerate import init_empty_weights
|
|
from transformers import AutoConfig, AutoProcessor, DiffusionGemmaForBlockDiffusion
|
|
|
|
load_start = time.perf_counter()
|
|
_dg_log(
|
|
"Preparing ModelOpt NVFP4 bridge. model_path=%s max_memory_gb=%s",
|
|
config.model_path,
|
|
config.max_memory_gb,
|
|
)
|
|
if 0 < float(config.max_memory_gb) < MIN_MODELOPT_NVFP4_BRIDGE_MEMORY_GB:
|
|
raise RuntimeError(
|
|
"The local ModelOpt NVFP4 bridge currently keeps DiffusionGemma weights on cuda:0 and does not "
|
|
f"support a {config.max_memory_gb:g} GiB memory cap. Use max_memory_gb >= "
|
|
f"{MIN_MODELOPT_NVFP4_BRIDGE_MEMORY_GB:g} for this bridge, or 0 for an explicit full-GPU load."
|
|
)
|
|
|
|
smoke = _nvfp4_bridge_smoke(config.model_path)
|
|
if not smoke.get("passed"):
|
|
raise RuntimeError(f"Comfy NVFP4 bridge smoke failed: {smoke.get('error', 'unknown error')}")
|
|
|
|
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
|
|
if device.type != "cuda":
|
|
raise RuntimeError("ModelOpt NVFP4 bridge requires CUDA.")
|
|
|
|
weight_map = _safetensors_index(config.model_path)
|
|
if not weight_map:
|
|
raise RuntimeError("No model.safetensors.index.json weight_map found for NVFP4 checkpoint.")
|
|
|
|
_dg_log("Loading DiffusionGemma config and processor.")
|
|
hf_config = AutoConfig.from_pretrained(config.model_path, local_files_only=config.local_files_only)
|
|
processor = AutoProcessor.from_pretrained(config.model_path, local_files_only=config.local_files_only)
|
|
with init_empty_weights():
|
|
model = DiffusionGemmaForBlockDiffusion(hf_config)
|
|
|
|
_replace_and_load_nvfp4_experts(model, hf_config, config.model_path, weight_map, device)
|
|
|
|
handles: dict[str, Any] = {}
|
|
try:
|
|
dense_keys = [key for key in sorted(weight_map) if not _is_modelopt_expert_tensor_key(key)]
|
|
_dg_log("Loading non-expert tensors: %s tensors.", len(dense_keys))
|
|
key_iter = dense_keys
|
|
tqdm_bar = None
|
|
if _progress_enabled():
|
|
try:
|
|
from tqdm.auto import tqdm
|
|
|
|
tqdm_bar = tqdm(dense_keys, total=len(dense_keys), desc="DiffusionGemma tensors", unit="tensor", leave=True)
|
|
key_iter = tqdm_bar
|
|
except Exception:
|
|
tqdm_bar = None
|
|
for key in key_iter:
|
|
if _is_modelopt_expert_tensor_key(key):
|
|
continue
|
|
tensor = _load_tensor_from_weight_map(config.model_path, weight_map, key, handles).to(device=device)
|
|
target_keys = [key]
|
|
tied_key = _decoder_to_encoder_language_key(key)
|
|
if tied_key and _model_has_leaf(model, tied_key):
|
|
target_keys.append(tied_key)
|
|
parent, leaf = _resolve_model_leaf(model, target_keys[0])
|
|
if leaf in getattr(parent, "_parameters", {}):
|
|
shared_value = torch.nn.Parameter(tensor, requires_grad=False)
|
|
else:
|
|
shared_value = tensor
|
|
for target_key in target_keys:
|
|
if _model_has_leaf(model, target_key):
|
|
_set_model_leaf(model, target_key, shared_value)
|
|
finally:
|
|
try:
|
|
if "tqdm_bar" in locals() and tqdm_bar is not None:
|
|
tqdm_bar.close()
|
|
except Exception:
|
|
pass
|
|
handles.clear()
|
|
|
|
if hasattr(model, "lm_head") and _model_has_leaf(model, "model.decoder.embed_tokens.weight"):
|
|
_set_model_leaf(model, "lm_head.weight", model.model.decoder.embed_tokens.weight)
|
|
if (Path(config.model_path) / "generation_config.json").exists():
|
|
model.generation_config = model.generation_config_class.from_pretrained(
|
|
config.model_path,
|
|
local_files_only=config.local_files_only,
|
|
)
|
|
model.eval()
|
|
model.to(device)
|
|
|
|
meta_names = [
|
|
name
|
|
for name, tensor in [*model.named_parameters(), *model.named_buffers()]
|
|
if getattr(tensor, "device", None) is not None and tensor.device.type == "meta"
|
|
]
|
|
if meta_names:
|
|
raise RuntimeError(f"NVFP4 bridge left {len(meta_names)} tensors on meta device, first: {meta_names[:8]}")
|
|
_dg_log("ModelOpt NVFP4 bridge ready in %.2fs.", time.perf_counter() - load_start)
|
|
return processor, model
|
|
|
|
|
|
def _load_transformers_model(config: RuntimeConfig) -> tuple[Any, Any]:
|
|
key = json.dumps(
|
|
{
|
|
"model_path": config.model_path,
|
|
"dtype": config.dtype,
|
|
"quantization": config.quantization,
|
|
"local_files_only": config.local_files_only,
|
|
"max_memory_gb": config.max_memory_gb,
|
|
},
|
|
sort_keys=True,
|
|
)
|
|
if key in _MODEL_CACHE:
|
|
_dg_log("Using cached Transformers runtime. backend=%s quantization=%s max_memory_gb=%s", config.backend, config.quantization, config.max_memory_gb)
|
|
return _MODEL_CACHE[key]
|
|
|
|
# Only one in-process DiffusionGemma profile can be resident safely on this
|
|
# workstation. Evict old max_memory/temperature/model profiles before a new
|
|
# load so changing loader settings does not double-allocate VRAM.
|
|
_release_transformers_runtime()
|
|
|
|
model_info = _model_path_info(config.model_path)
|
|
try:
|
|
_dg_log(
|
|
"Loading Transformers runtime. path_kind=%s quantization=%s max_memory_gb=%s",
|
|
model_info.get("path_kind", "unknown"),
|
|
config.quantization,
|
|
config.max_memory_gb,
|
|
)
|
|
if model_info["path_kind"] == "nvfp4_hf_repo" and model_info["quant_method"] == "modelopt":
|
|
if config.quantization == "modelopt_nvfp4":
|
|
processor, model = _load_modelopt_nvfp4_bridge(config)
|
|
_MODEL_CACHE[key] = (processor, model)
|
|
_dg_log("Transformers runtime cached.")
|
|
return processor, model
|
|
allow_unsupported = os.environ.get("DG_ALLOW_UNSUPPORTED_MODELOPT_NVFP4_LOAD", "").strip() == "1"
|
|
if not allow_unsupported:
|
|
raise RuntimeError(
|
|
"The selected NVFP4 checkpoint uses quant_method=modelopt. This ComfyUI runtime has "
|
|
"DiffusionGemmaForBlockDiffusion, but Transformers does not support modelopt NVFP4 "
|
|
"deserialization here. Choose quantization=modelopt_nvfp4 to use the Comfy NVFP4 bridge, "
|
|
"or set DG_ALLOW_UNSUPPORTED_MODELOPT_NVFP4_LOAD=1 only for future runtime experiments."
|
|
)
|
|
|
|
from transformers import AutoProcessor, DiffusionGemmaForBlockDiffusion
|
|
|
|
kwargs: dict[str, Any] = {
|
|
"dtype": "auto" if config.dtype == "auto" else config.dtype,
|
|
"device_map": "auto",
|
|
"local_files_only": config.local_files_only,
|
|
}
|
|
if config.max_memory_gb > 0:
|
|
kwargs["max_memory"] = {0: f"{config.max_memory_gb}GiB", "cpu": "64GiB"}
|
|
|
|
processor = AutoProcessor.from_pretrained(config.model_path, local_files_only=config.local_files_only)
|
|
model = DiffusionGemmaForBlockDiffusion.from_pretrained(config.model_path, **kwargs)
|
|
_MODEL_CACHE[key] = (processor, model)
|
|
_dg_log("Transformers runtime cached.")
|
|
return processor, model
|
|
except Exception:
|
|
_release_transformers_runtime()
|
|
raise
|
|
|
|
|
|
def _transformers_messages_and_processor_kwargs(
|
|
prompt: str,
|
|
media_context: MediaContext | None,
|
|
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
|
|
content: list[dict[str, Any]] = []
|
|
pil_images = _image_batch_to_pil(media_context.images if isinstance(media_context, MediaContext) else None)
|
|
media_metadata = media_context.metadata if isinstance(media_context, MediaContext) else {}
|
|
source = media_metadata.get("source", "none") if isinstance(media_metadata, dict) else "none"
|
|
processor_kwargs: dict[str, Any] = {}
|
|
if pil_images and source in {"video", "image+video"}:
|
|
frame_count = len(pil_images)
|
|
transport = str(media_metadata.get("transformers_video_transport", _video_transport_mode()))
|
|
if _is_image_identity_video_control(media_metadata):
|
|
transport = "sampled_frame_images"
|
|
if transport == "video_tokens":
|
|
content.append({"type": "video", "video": pil_images})
|
|
duration = _safe_float(media_metadata.get("duration_seconds"), 0.0) if isinstance(media_metadata, dict) else 0.0
|
|
effective_fps = frame_count / duration if duration > 0 else _safe_float(media_metadata.get("sample_fps"), 1.0)
|
|
video_metadata = {
|
|
"total_num_frames": frame_count,
|
|
"fps": max(0.001, effective_fps),
|
|
"width": int(media_metadata.get("width", 0) or 0) or None,
|
|
"height": int(media_metadata.get("height", 0) or 0) or None,
|
|
"duration": duration or None,
|
|
}
|
|
processor_kwargs["videos_kwargs"] = {
|
|
"num_frames": max(1, frame_count),
|
|
"video_metadata": video_metadata,
|
|
}
|
|
else:
|
|
content.append({"type": "text", "text": _sampled_video_frame_sequence_intro(media_metadata, frame_count)})
|
|
if _is_image_identity_video_control(media_metadata):
|
|
reference_count = int(max(1.0, _safe_float(media_metadata.get("reference_image_count"), 1.0)))
|
|
reference_count = min(reference_count, frame_count)
|
|
reference_images = pil_images[:reference_count]
|
|
video_images = pil_images[reference_count:]
|
|
content.append(
|
|
{
|
|
"type": "text",
|
|
"text": (
|
|
"IDENTITY REFERENCE IMAGE(S): use only these image(s) for the output subject's face, hair, "
|
|
"body type, wardrobe, accessories, styling, and distinguishing visual traits. These appearance "
|
|
"traits override the control-video subject."
|
|
),
|
|
}
|
|
)
|
|
for image in reference_images:
|
|
content.append({"type": "image", "image": image})
|
|
content.append(
|
|
{
|
|
"type": "text",
|
|
"text": (
|
|
"VIDEO CONTROL FRAMES: use the following frames only for pose, action, blocking, depth/canny/edge "
|
|
"layout, camera motion, timing, composition, lighting/capture, and scene geometry. Do not copy the "
|
|
"control-video subject's clothing, hair, face, body type, accessories, or styling."
|
|
),
|
|
}
|
|
)
|
|
for image in video_images:
|
|
content.append({"type": "image", "image": image})
|
|
content.append(
|
|
{
|
|
"type": "text",
|
|
"text": (
|
|
"End video control frames. The identity reference is not part of the video timeline. The final prompt "
|
|
"must describe the reference-image subject, including reference wardrobe and accessories, performing "
|
|
"the video motion/control structure."
|
|
),
|
|
}
|
|
)
|
|
else:
|
|
for image in pil_images:
|
|
content.append({"type": "image", "image": image})
|
|
else:
|
|
for image in pil_images:
|
|
content.append({"type": "image", "image": image})
|
|
content.append({"type": "text", "text": prompt})
|
|
return [{"role": "user", "content": content}], processor_kwargs
|
|
|
|
|
|
def _run_transformers(
|
|
config: RuntimeConfig,
|
|
prompt: str,
|
|
media_context: MediaContext | None,
|
|
max_new_tokens: int,
|
|
node_id: str | None = None,
|
|
) -> str:
|
|
run_start = time.perf_counter()
|
|
processor, model = _load_transformers_model(config)
|
|
messages, processor_kwargs = _transformers_messages_and_processor_kwargs(prompt, media_context)
|
|
content = messages[0].get("content", []) if messages else []
|
|
image_count = sum(1 for item in content if isinstance(item, dict) and item.get("type") == "image")
|
|
video_count = sum(1 for item in content if isinstance(item, dict) and item.get("type") == "video")
|
|
_dg_log(
|
|
"Preparing processor inputs. images=%s videos=%s max_new_tokens=%s temperature=%.3g",
|
|
image_count,
|
|
video_count,
|
|
max_new_tokens,
|
|
config.temperature,
|
|
)
|
|
|
|
processor_start = time.perf_counter()
|
|
inputs = processor.apply_chat_template(
|
|
messages,
|
|
tokenize=True,
|
|
return_dict=True,
|
|
return_tensors="pt",
|
|
add_generation_prompt=True,
|
|
processor_kwargs=processor_kwargs,
|
|
)
|
|
input_tokens = int(inputs["input_ids"].shape[-1]) if "input_ids" in inputs else 0
|
|
_dg_log("Processor inputs ready in %.2fs. input_tokens=%s", time.perf_counter() - processor_start, input_tokens)
|
|
try:
|
|
import torch
|
|
|
|
inputs = {
|
|
key: value.to(model.device) if isinstance(value, torch.Tensor) and hasattr(model, "device") else value
|
|
for key, value in inputs.items()
|
|
}
|
|
except Exception:
|
|
pass
|
|
|
|
generate_kwargs: dict[str, Any] = {"max_new_tokens": max_new_tokens}
|
|
denoise_steps = os.environ.get("DG_MAX_DENOISING_STEPS", "").strip()
|
|
if denoise_steps:
|
|
try:
|
|
generate_kwargs["max_denoising_steps"] = max(1, int(denoise_steps))
|
|
except Exception:
|
|
pass
|
|
if config.temperature > 0:
|
|
generate_kwargs["do_sample"] = True
|
|
generate_kwargs["temperature"] = config.temperature
|
|
_dg_log("Generation started. max_new_tokens=%s", max_new_tokens)
|
|
|
|
def do_generate() -> Any:
|
|
import torch
|
|
|
|
with torch.inference_mode():
|
|
return model.generate(**inputs, **generate_kwargs)
|
|
|
|
generate_start = time.perf_counter()
|
|
generated = _run_blocking_with_progress(
|
|
"DiffusionGemma generate",
|
|
do_generate,
|
|
node_id=node_id,
|
|
estimated_seconds=max(30.0, min(300.0, 45.0 + (float(max_new_tokens) * 0.12))),
|
|
)
|
|
_dg_log("Generation finished in %.2fs.", time.perf_counter() - generate_start)
|
|
sequences = getattr(generated, "sequences", generated)
|
|
decode_start = time.perf_counter()
|
|
if hasattr(processor, "batch_decode"):
|
|
input_len = inputs["input_ids"].shape[-1] if "input_ids" in inputs else 0
|
|
try:
|
|
text = processor.batch_decode(sequences[:, input_len:], skip_special_tokens=True)[0]
|
|
except Exception:
|
|
text = processor.batch_decode(sequences, skip_special_tokens=True)[0]
|
|
_dg_log("Decode finished in %.2fs. output_chars=%s total_runtime=%.2fs", time.perf_counter() - decode_start, len(text), time.perf_counter() - run_start)
|
|
return text
|
|
if hasattr(processor, "decode"):
|
|
text = processor.decode(sequences[0], skip_special_tokens=True)
|
|
_dg_log("Decode finished in %.2fs. output_chars=%s total_runtime=%.2fs", time.perf_counter() - decode_start, len(text), time.perf_counter() - run_start)
|
|
return text
|
|
text = str(sequences)
|
|
_dg_log("Decode finished in %.2fs. output_chars=%s total_runtime=%.2fs", time.perf_counter() - decode_start, len(text), time.perf_counter() - run_start)
|
|
return text
|
|
|
|
|
|
def _run_gguf_subprocess(config: RuntimeConfig, prompt: str, max_new_tokens: int, node_id: str | None = None) -> str:
|
|
if not config.cli_path or not Path(config.cli_path).exists():
|
|
raise RuntimeError("GGUF backend requires a valid llama-diffusion-cli path.")
|
|
if not Path(config.model_path).exists():
|
|
raise RuntimeError("GGUF backend requires model_path to point at a .gguf file.")
|
|
cmd = [config.cli_path, "-m", config.model_path, "-p", prompt, "-n", str(max_new_tokens)]
|
|
cmd.extend(_split_extra_args(config.extra_args))
|
|
creationflags = subprocess.CREATE_NO_WINDOW if os.name == "nt" and hasattr(subprocess, "CREATE_NO_WINDOW") else 0
|
|
_dg_log("GGUF subprocess started. max_new_tokens=%s", max_new_tokens)
|
|
|
|
def run_subprocess() -> Any:
|
|
return subprocess.run(
|
|
cmd,
|
|
capture_output=True,
|
|
text=True,
|
|
check=False,
|
|
env=_subprocess_env_for_cli(config.cli_path),
|
|
creationflags=creationflags,
|
|
timeout=900,
|
|
)
|
|
|
|
subprocess_start = time.perf_counter()
|
|
result = _run_blocking_with_progress(
|
|
"DiffusionGemma GGUF",
|
|
run_subprocess,
|
|
node_id=node_id,
|
|
estimated_seconds=max(30.0, min(300.0, 45.0 + (float(max_new_tokens) * 0.12))),
|
|
)
|
|
if result.returncode != 0:
|
|
raise RuntimeError(f"GGUF subprocess failed with code {result.returncode}: {result.stderr.strip()}")
|
|
_dg_log("GGUF subprocess finished in %.2fs.", time.perf_counter() - subprocess_start)
|
|
return _strip_cli_runtime_lines(result.stdout)
|
|
|
|
|
|
def _run_backend(
|
|
config: RuntimeConfig,
|
|
prompt: str,
|
|
media_context: MediaContext | None,
|
|
max_new_tokens: int,
|
|
node_id: str | None = None,
|
|
) -> str:
|
|
if config.backend == "transformers_inprocess":
|
|
return _run_transformers(config, prompt, media_context, max_new_tokens, node_id=node_id)
|
|
if config.backend == "gguf_subprocess":
|
|
return _run_gguf_subprocess(config, prompt, max_new_tokens, node_id=node_id)
|
|
if config.backend in {"qwen_vl_placeholder", "minicpm_v_placeholder"}:
|
|
raise RuntimeError(
|
|
f"{config.backend} is a reserved benchmark lane. Install or wire a concrete local VLM backend before enabling runtime_required."
|
|
)
|
|
raise RuntimeError(f"Backend {config.backend!r} does not provide model inference.")
|
|
|
|
|
|
def _maybe_unload(config: RuntimeConfig) -> None:
|
|
if config.unload_policy != "unload_after_run":
|
|
return
|
|
_release_transformers_runtime()
|
|
|
|
|
|
def _runtime_config_with_temperature(config: RuntimeConfig, temperature: float) -> RuntimeConfig:
|
|
temp = max(0.0, min(2.0, _safe_float(temperature, config.temperature)))
|
|
updated = RuntimeConfig(
|
|
model_path=config.model_path,
|
|
backend=config.backend,
|
|
dtype=config.dtype,
|
|
quantization=config.quantization,
|
|
local_files_only=config.local_files_only,
|
|
unload_policy=config.unload_policy,
|
|
max_memory_gb=config.max_memory_gb,
|
|
cli_path=config.cli_path,
|
|
extra_args=_extra_args_with_temperature(config.extra_args, temp),
|
|
temperature=temp,
|
|
fallback_backend=config.fallback_backend,
|
|
status=config.status.copy(),
|
|
)
|
|
updated.status["temperature"] = temp
|
|
return updated
|
|
|
|
|
|
def _run_generation_packet(
|
|
model_config: RuntimeConfig,
|
|
gemma_context: GemmaContext,
|
|
target_config: TargetProfileConfig | dict[str, Any],
|
|
master_prompt: str = DEFAULT_MASTER_PROMPT,
|
|
runtime_required: bool = False,
|
|
max_new_tokens: int = 768,
|
|
max_output_chars: int = 8000,
|
|
creativity_mode: str = "editorial",
|
|
creative_strength: float = 0.6,
|
|
thinking_mode: str = "off",
|
|
node_id: str | None = None,
|
|
) -> tuple[dict[str, Any], str, str, str, bool, str, dict[str, Any]]:
|
|
total_start = time.perf_counter()
|
|
if not isinstance(model_config, RuntimeConfig):
|
|
raise TypeError("model_config must come from DiffusionGemma Model Loader.")
|
|
context = gemma_context if isinstance(gemma_context, GemmaContext) else _gemma_context_from_media("")
|
|
target = _target_profile_config_to_dict(target_config)
|
|
creativity_mode = _normalize_creativity_mode(creativity_mode)
|
|
creative_strength = _normalize_creative_strength(creative_strength)
|
|
thinking_mode = _normalize_thinking_mode(thinking_mode)
|
|
media_notes = _media_policy_notes(model_config, context)
|
|
media_metadata = context.media_metadata.copy()
|
|
pixels_present = bool(context.images is not None)
|
|
pixels_sent = bool(pixels_present and _runtime_supports_pixels(model_config))
|
|
visual_description = _clean_visual_description(context.visual_description or str(media_metadata.get("visual_description", "")))
|
|
media_metadata["pixel_tensor_present"] = pixels_present
|
|
media_metadata["pixels_sent_to_backend"] = pixels_sent
|
|
if pixels_sent and str(media_metadata.get("source", "none")).lower() in {"video", "image+video"}:
|
|
media_metadata["transformers_video_transport"] = (
|
|
"sampled_frame_images" if _is_image_identity_video_control(media_metadata) else _video_transport_mode()
|
|
)
|
|
if visual_description:
|
|
media_metadata["visual_description"] = visual_description
|
|
media_metadata.setdefault("visual_description_source", "upstream_text")
|
|
media_metadata["visual_grounding_mode"] = _media_grounding_mode(media_metadata)
|
|
warnings = list(media_metadata.get("warnings", [])) if isinstance(media_metadata.get("warnings"), list) else []
|
|
for note in media_notes:
|
|
if note not in warnings:
|
|
warnings.append(note)
|
|
if pixels_present and not pixels_sent and not visual_description:
|
|
note = "No visual_description was supplied; the active backend cannot inspect image/video pixels and will use only text plus metadata."
|
|
if note not in warnings:
|
|
warnings.append(note)
|
|
media_metadata["warnings"] = warnings
|
|
working_context = GemmaContext(
|
|
user_prompt=context.user_prompt,
|
|
images=context.images,
|
|
source=context.source,
|
|
media_metadata=media_metadata,
|
|
visual_description=visual_description,
|
|
warnings=warnings,
|
|
)
|
|
model_prompt = _build_model_prompt(
|
|
working_context.user_prompt,
|
|
master_prompt,
|
|
target["target_profile"],
|
|
media_metadata,
|
|
target["audio_mode"],
|
|
target["audio_guidance"],
|
|
target["target_duration_seconds"],
|
|
target["ltx_style"],
|
|
target["ideogram_aspect_ratio"],
|
|
target["ideogram_render_style"],
|
|
target["ideogram_exact_text"],
|
|
target["ideogram_json_output"],
|
|
creativity_mode,
|
|
creative_strength,
|
|
target["negative_prompt_mode"],
|
|
target["negative_prompt_guidance"],
|
|
thinking_mode,
|
|
)
|
|
raw_output = ""
|
|
answer_text = ""
|
|
reasoning_text = ""
|
|
fallback_reason = ""
|
|
packet: dict[str, Any]
|
|
parse_valid = False
|
|
salvage_warning = ""
|
|
backend_seconds = 0.0
|
|
postprocess_seconds = 0.0
|
|
try:
|
|
if model_config.backend == "template":
|
|
raise RuntimeError("Template backend selected.")
|
|
backend_media_context = _media_context_from_gemma_context(working_context) if _runtime_supports_pixels(model_config) else None
|
|
backend_start = time.perf_counter()
|
|
_dg_log("Backend dispatch. backend=%s pixels_sent=%s", model_config.backend, bool(backend_media_context))
|
|
raw_output = _run_backend(model_config, model_prompt, backend_media_context, int(max_new_tokens), node_id=node_id)
|
|
backend_seconds = time.perf_counter() - backend_start
|
|
_dg_log("Backend returned in %.2fs. raw_output_chars=%s", backend_seconds, len(raw_output))
|
|
postprocess_start = time.perf_counter()
|
|
reasoning_text, answer_text = _split_reasoning_and_answer(raw_output)
|
|
packet, parse_valid, salvage_warning = repair_or_salvage_prompt_packet(
|
|
answer_text or raw_output,
|
|
working_context,
|
|
target,
|
|
int(max_output_chars),
|
|
creativity_mode,
|
|
creative_strength,
|
|
)
|
|
postprocess_seconds = time.perf_counter() - postprocess_start
|
|
except Exception as exc:
|
|
fallback_reason = str(exc)
|
|
if model_config.backend == "transformers_inprocess":
|
|
_release_transformers_runtime()
|
|
if runtime_required:
|
|
raise RuntimeError(f"DiffusionGemma runtime failed and runtime_required is true: {exc}") from exc
|
|
postprocess_start = time.perf_counter()
|
|
packet = _template_packet(
|
|
working_context.user_prompt,
|
|
target["target_profile"],
|
|
media_metadata,
|
|
_safe_float(media_metadata.get("duration_seconds"), 0.0),
|
|
target["audio_mode"],
|
|
target["audio_guidance"],
|
|
target["target_duration_seconds"],
|
|
target["ltx_style"],
|
|
target["ideogram_aspect_ratio"],
|
|
target["ideogram_render_style"],
|
|
target["ideogram_exact_text"],
|
|
target["ideogram_json_output"],
|
|
creativity_mode,
|
|
creative_strength,
|
|
target["negative_prompt_mode"],
|
|
target["negative_prompt_guidance"],
|
|
)
|
|
packet.setdefault("metadata", {})["fallback_reason"] = fallback_reason
|
|
postprocess_seconds = time.perf_counter() - postprocess_start
|
|
finally:
|
|
_maybe_unload(model_config)
|
|
|
|
total_seconds = time.perf_counter() - total_start
|
|
final_json_preview = _json_dumps(packet)
|
|
metadata = packet.get("metadata") if isinstance(packet.get("metadata"), dict) else {}
|
|
metadata.update(
|
|
{
|
|
"backend": model_config.backend,
|
|
"runtime_ready": bool(model_config.status.get("ready")),
|
|
"raw_output_available": bool(raw_output),
|
|
"used_template_fallback": bool(fallback_reason),
|
|
"json_parse_valid": bool(parse_valid),
|
|
"salvage_warning": salvage_warning,
|
|
"media_policy": model_config.status.get("media_policy", "unknown"),
|
|
"media": media_metadata,
|
|
"audio_mode": target["audio_mode"],
|
|
"audio_guidance": target["audio_guidance"],
|
|
"target_duration_seconds": _effective_duration_seconds(media_metadata, target["target_duration_seconds"]),
|
|
"ltx_style": target["ltx_style"],
|
|
"ideogram_aspect_ratio": target["ideogram_aspect_ratio"],
|
|
"ideogram_render_style": target["ideogram_render_style"],
|
|
"ideogram_exact_text": target["ideogram_exact_text"],
|
|
"ideogram_json_output": bool(target["ideogram_json_output"]),
|
|
"creativity_mode": creativity_mode,
|
|
"creative_strength": creative_strength,
|
|
"thinking_mode": thinking_mode,
|
|
"reasoning_available": bool(reasoning_text),
|
|
"negative_prompt_mode": target["negative_prompt_mode"],
|
|
"negative_prompt_guidance": target["negative_prompt_guidance"],
|
|
"timing": {
|
|
"backend_generation_seconds": round(backend_seconds, 3),
|
|
"postprocess_seconds": round(postprocess_seconds, 3),
|
|
"total_node_seconds": round(total_seconds, 3),
|
|
},
|
|
"lengths": {
|
|
"model_prompt_chars": len(model_prompt),
|
|
"raw_output_chars": len(raw_output),
|
|
"answer_text_chars": len(answer_text),
|
|
"reasoning_text_chars": len(reasoning_text),
|
|
"final_json_chars": len(final_json_preview),
|
|
"max_new_tokens": int(max_new_tokens),
|
|
},
|
|
}
|
|
)
|
|
_copy_synthesis_metadata(metadata, media_metadata)
|
|
readiness_prompt = ""
|
|
if target["target_profile"] == "ideogram4":
|
|
readiness_prompt = _ideogram_prompt_value_to_text(
|
|
packet.get("ideogram_prompt", ""),
|
|
_clean_prompt_request_text(working_context.user_prompt),
|
|
target["ideogram_aspect_ratio"],
|
|
target["ideogram_render_style"],
|
|
target["ideogram_exact_text"],
|
|
target["ideogram_json_output"],
|
|
int(max_output_chars),
|
|
creativity_mode,
|
|
creative_strength,
|
|
)
|
|
else:
|
|
readiness_prompt = _prompt_field_to_text(
|
|
packet.get("ltx_prompt", ""),
|
|
_clean_prompt_request_text(working_context.user_prompt),
|
|
int(max_output_chars),
|
|
)
|
|
ready_for_generation, blocked_reasons = _packet_generation_readiness(
|
|
packet,
|
|
working_context,
|
|
target,
|
|
readiness_prompt,
|
|
media_metadata,
|
|
)
|
|
metadata["ready_for_generation"] = bool(ready_for_generation)
|
|
metadata["blocked_reasons"] = blocked_reasons
|
|
metadata["claim_verification_passed"] = not any(
|
|
reason.startswith("metadata_only_") or reason.startswith("visual_description_contradiction")
|
|
for reason in blocked_reasons
|
|
)
|
|
packet["metadata"] = metadata
|
|
return packet, raw_output, reasoning_text, fallback_reason, parse_valid, salvage_warning, metadata
|
|
|
|
|
|
class DiffusionGemmaModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_path": (
|
|
"STRING",
|
|
{
|
|
"default": DEFAULT_ADVANCED_MODEL_PATH,
|
|
"multiline": False,
|
|
"tooltip": "Whole Hugging Face repo folder for the in-process Transformers/NVFP4 runtime.",
|
|
},
|
|
),
|
|
"backend": (
|
|
[
|
|
"transformers_inprocess",
|
|
],
|
|
{"default": "transformers_inprocess"},
|
|
),
|
|
"dtype": (["auto", "bfloat16", "float16", "float32"], {"default": "auto"}),
|
|
"quantization": (
|
|
["modelopt_nvfp4", "none", "bitsandbytes_4bit", "quanto", "nvfp4_metadata_only"],
|
|
{"default": "modelopt_nvfp4"},
|
|
),
|
|
"local_files_only": ("BOOLEAN", {"default": True}),
|
|
"unload_policy": (["keep_loaded", "unload_after_run"], {"default": "keep_loaded"}),
|
|
"max_memory_gb": (
|
|
MEMORY_PRESET_CHOICES,
|
|
{
|
|
"default": MEMORY_PRESET_SAFE_20,
|
|
"tooltip": "NVFP4 bridge memory profile. 20 GB is the safest practical default; 18 GB is the supported floor; 0 requests full GPU load.",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"temperature": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.45,
|
|
"min": 0.0,
|
|
"max": 2.0,
|
|
"step": 0.05,
|
|
"tooltip": "Generation sampling temperature. 0.2 is conservative; 0.45-0.8 is usually more creative.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (MODEL_TYPE, "STRING")
|
|
RETURN_NAMES = ("model_config", "runtime_status_json")
|
|
FUNCTION = "load"
|
|
CATEGORY = CATEGORY
|
|
|
|
def load(
|
|
self,
|
|
model_path: str,
|
|
backend: str = "transformers_inprocess",
|
|
dtype: str = "auto",
|
|
quantization: str = "modelopt_nvfp4",
|
|
local_files_only: bool = True,
|
|
unload_policy: str = "keep_loaded",
|
|
max_memory_gb: Any = MEMORY_PRESET_SAFE_20,
|
|
cli_path: str = "",
|
|
extra_args: str = "",
|
|
temperature: float = 0.45,
|
|
):
|
|
valid_backends = {"transformers_inprocess", "gguf_subprocess", "qwen_vl_placeholder", "minicpm_v_placeholder", "template"}
|
|
valid_dtypes = {"auto", "bfloat16", "float16", "float32"}
|
|
valid_quantization = {"none", "modelopt_nvfp4", "bitsandbytes_4bit", "quanto", "nvfp4_metadata_only", "gguf_q4_q5"}
|
|
valid_unload = {"keep_loaded", "unload_after_run"}
|
|
backend = backend if backend in valid_backends else ("gguf_subprocess" if Path(DEFAULT_GGUF_MODEL_PATH).exists() else "transformers_inprocess")
|
|
dtype = dtype if dtype in valid_dtypes else "auto"
|
|
quantization = quantization if quantization in valid_quantization else "none"
|
|
unload_policy = unload_policy if unload_policy in valid_unload else "keep_loaded"
|
|
cli_path = (cli_path or "").strip()
|
|
extra_args = (extra_args or "").strip()
|
|
if cli_path.startswith("-") or (cli_path and not Path(cli_path).exists() and Path(DEFAULT_GGUF_CLI_PATH).exists()):
|
|
if not extra_args or extra_args == cli_path:
|
|
extra_args = cli_path if cli_path.startswith("-") else extra_args
|
|
cli_path = DEFAULT_GGUF_CLI_PATH
|
|
if not cli_path and backend == "gguf_subprocess":
|
|
cli_path = DEFAULT_GGUF_CLI_PATH
|
|
if not extra_args and backend == "gguf_subprocess":
|
|
extra_args = DEFAULT_GGUF_EXTRA_ARGS
|
|
temp_value = max(0.0, min(2.0, _safe_float(temperature, 0.45)))
|
|
memory_gb = _memory_preset_to_gb(max_memory_gb)
|
|
effective_extra_args = _extra_args_with_temperature(extra_args, temp_value)
|
|
config = RuntimeConfig(
|
|
model_path=model_path.strip(),
|
|
backend=backend,
|
|
dtype=dtype,
|
|
quantization=quantization,
|
|
local_files_only=bool(local_files_only),
|
|
unload_policy=unload_policy,
|
|
max_memory_gb=memory_gb,
|
|
cli_path=cli_path,
|
|
extra_args=effective_extra_args,
|
|
temperature=temp_value,
|
|
)
|
|
config.status = _runtime_status(config)
|
|
return (config, _json_dumps(config.status))
|
|
|
|
|
|
class DiffusionGemmaSimpleModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_path": (
|
|
"STRING",
|
|
{
|
|
"default": DEFAULT_MODEL_PATH,
|
|
"multiline": False,
|
|
"tooltip": "Path to the local DiffusionGemma GGUF file.",
|
|
},
|
|
),
|
|
"temperature": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.45,
|
|
"min": 0.0,
|
|
"max": 2.0,
|
|
"step": 0.05,
|
|
"tooltip": "Generation randomness. 0.2 is literal; 0.45 balanced; 0.65+ is more creative.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (MODEL_TYPE, "STRING")
|
|
RETURN_NAMES = ("model_config", "runtime_status_json")
|
|
FUNCTION = "load"
|
|
CATEGORY = COMPATIBILITY_CATEGORY
|
|
|
|
def load(self, model_path: str, temperature: float = 0.45):
|
|
return DiffusionGemmaModelLoader().load(
|
|
model_path,
|
|
"gguf_subprocess",
|
|
"auto",
|
|
"none",
|
|
True,
|
|
"keep_loaded",
|
|
0.0,
|
|
DEFAULT_GGUF_CLI_PATH,
|
|
DEFAULT_GGUF_EXTRA_ARGS,
|
|
temperature,
|
|
)
|
|
|
|
|
|
class DiffusionGemmaModelPathLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_path": (
|
|
"STRING",
|
|
{
|
|
"default": DEFAULT_MODEL_PATH,
|
|
"multiline": False,
|
|
"tooltip": "Path to the local DiffusionGemma GGUF file. Sampling temperature is set on DiffusionGemma CoT Generator.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (MODEL_TYPE, "STRING")
|
|
RETURN_NAMES = ("model_config", "runtime_status_json")
|
|
FUNCTION = "load"
|
|
CATEGORY = COMPATIBILITY_CATEGORY
|
|
|
|
def load(self, model_path: str):
|
|
return DiffusionGemmaModelLoader().load(
|
|
model_path,
|
|
"gguf_subprocess",
|
|
"auto",
|
|
"none",
|
|
True,
|
|
"keep_loaded",
|
|
0.0,
|
|
DEFAULT_GGUF_CLI_PATH,
|
|
DEFAULT_GGUF_EXTRA_ARGS,
|
|
0.45,
|
|
)
|
|
|
|
|
|
class DiffusionGemmaMediaSampler:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"sample_fps": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 24.0, "step": 0.1}),
|
|
"max_duration_seconds": ("FLOAT", {"default": 60.0, "min": 1.0, "max": 600.0, "step": 1.0}),
|
|
"max_frames": ("INT", {"default": 60, "min": 1, "max": 512}),
|
|
"overlong_policy": (["trim", "reject"], {"default": "trim"}),
|
|
"include_image_with_video": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE",),
|
|
"video": ("VIDEO",),
|
|
"visual_description": ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True}),
|
|
"media_synthesis_mode": (MEDIA_SYNTHESIS_MODE_CHOICES, {"default": "video_recreation"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (MEDIA_TYPE, "STRING")
|
|
RETURN_NAMES = ("media_context", "media_metadata_json")
|
|
FUNCTION = "sample"
|
|
CATEGORY = OPTIONAL_CATEGORY
|
|
|
|
def sample(
|
|
self,
|
|
sample_fps: float,
|
|
max_duration_seconds: float,
|
|
max_frames: int,
|
|
overlong_policy: str,
|
|
include_image_with_video: bool,
|
|
image=None,
|
|
video=None,
|
|
visual_description: str = "",
|
|
media_synthesis_mode: str = "video_recreation",
|
|
):
|
|
synthesis_mode = _normalize_media_synthesis_mode(media_synthesis_mode)
|
|
metadata: dict[str, Any] = {
|
|
"source": "none",
|
|
"sample_fps": float(sample_fps),
|
|
"max_duration_seconds": float(max_duration_seconds),
|
|
"max_frames": int(max_frames),
|
|
"media_synthesis_mode": synthesis_mode,
|
|
"image_role": "none",
|
|
"video_role": "none",
|
|
"reference_image_count": 0,
|
|
"reference_image_backend_attached": False,
|
|
"video_sampled_frame_count": 0,
|
|
"synthesis_contract": _synthesis_contract_text({"media_synthesis_mode": synthesis_mode}),
|
|
"warnings": [],
|
|
}
|
|
sampled = None
|
|
|
|
if video is not None:
|
|
duration = _safe_float(video.get_duration(), 0.0)
|
|
metadata["source"] = "video"
|
|
metadata["duration_seconds"] = duration
|
|
try:
|
|
metadata["frame_count"] = int(video.get_frame_count())
|
|
except Exception as exc:
|
|
metadata["warnings"].append(f"Could not read frame count without materializing video: {exc}")
|
|
try:
|
|
width, height = video.get_dimensions()
|
|
metadata["width"] = int(width)
|
|
metadata["height"] = int(height)
|
|
except Exception as exc:
|
|
metadata["warnings"].append(f"Could not read video dimensions: {exc}")
|
|
|
|
source_video = video
|
|
if duration > max_duration_seconds:
|
|
if overlong_policy == "reject":
|
|
raise ValueError(
|
|
f"Video is {duration:.2f}s, above the {max_duration_seconds:.2f}s DiffusionGemma analysis cap."
|
|
)
|
|
source_video = video.as_trimmed(0.0, float(max_duration_seconds), strict_duration=False)
|
|
metadata["trimmed_to_seconds"] = float(max_duration_seconds)
|
|
metadata["warnings"].append("Video was trimmed for DiffusionGemma analysis.")
|
|
if source_video is None:
|
|
raise ValueError("Video trim failed.")
|
|
|
|
components = source_video.get_components()
|
|
frames = components.images
|
|
frame_count = int(frames.shape[0])
|
|
fps = float(components.frame_rate)
|
|
stride = max(1, int(round(fps / max(sample_fps, 0.1))))
|
|
indices = list(range(0, frame_count, stride))
|
|
if len(indices) > max_frames:
|
|
step = max(1, math.ceil(len(indices) / max_frames))
|
|
indices = indices[::step][:max_frames]
|
|
try:
|
|
import torch
|
|
|
|
index_tensor = torch.tensor(indices, dtype=torch.long, device=frames.device)
|
|
sampled = frames.index_select(0, index_tensor)
|
|
except Exception:
|
|
sampled = frames[indices]
|
|
metadata["source_fps"] = fps
|
|
metadata["sampled_indices"] = indices
|
|
metadata["sampled_frame_count"] = len(indices)
|
|
metadata["video_sampled_frame_count"] = len(indices)
|
|
metadata["video_role"] = (
|
|
"control_structure_pose_depth_canny_composition_motion_camera"
|
|
if synthesis_mode == "image_identity_video_control"
|
|
else "primary_video_recreation_source"
|
|
)
|
|
|
|
if image is not None:
|
|
metadata["reference_image_count"] = 1
|
|
metadata["image_role"] = (
|
|
"identity_reference"
|
|
if synthesis_mode == "image_identity_video_control"
|
|
else ("optional_still_reference" if sampled is not None else "primary_image_source")
|
|
)
|
|
should_include_image = bool(include_image_with_video or (synthesis_mode == "image_identity_video_control" and sampled is not None))
|
|
if sampled is not None and should_include_image:
|
|
try:
|
|
import torch
|
|
|
|
reference_for_backend, reference_report = _fit_reference_image_to_video_frames(image, sampled)
|
|
metadata.update(reference_report)
|
|
sampled = torch.cat([reference_for_backend, sampled], dim=0)
|
|
metadata["source"] = "image+video"
|
|
metadata["sampled_frame_count"] = int(sampled.shape[0])
|
|
metadata["reference_image_backend_attached"] = True
|
|
if metadata.get("reference_image_resized_for_video_control"):
|
|
metadata["warnings"].append(
|
|
"Reference image was aspect-fitted and padded to match sampled video frame dimensions for synthesis evidence."
|
|
)
|
|
except Exception as exc:
|
|
metadata["warnings"].append(f"Could not prepend image to sampled video frames: {exc}")
|
|
elif sampled is None:
|
|
sampled = image
|
|
metadata["source"] = "image"
|
|
metadata["sampled_frame_count"] = int(image.shape[0]) if hasattr(image, "shape") else 1
|
|
metadata["reference_image_backend_attached"] = True
|
|
elif synthesis_mode == "image_identity_video_control":
|
|
metadata["image_role"] = "missing_identity_reference"
|
|
|
|
if video is None and synthesis_mode == "image_identity_video_control":
|
|
metadata["video_role"] = "missing_control_video"
|
|
|
|
cleaned_visual_description = _clean_visual_description(visual_description)
|
|
if cleaned_visual_description:
|
|
metadata["visual_description"] = cleaned_visual_description
|
|
metadata["visual_description_source"] = "upstream_text"
|
|
|
|
context = MediaContext(images=sampled, source=metadata["source"], metadata=metadata)
|
|
return (context, _json_dumps(metadata))
|
|
|
|
|
|
class DiffusionGemmaContextHub:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"user_prompt": (
|
|
"STRING",
|
|
{"default": "", "multiline": True, "forceInput": True, "dynamicPrompts": True},
|
|
),
|
|
"sample_fps": ("FLOAT", {"default": 1.0, "min": 0.1, "max": 24.0, "step": 0.1}),
|
|
"max_duration_seconds": ("FLOAT", {"default": 60.0, "min": 1.0, "max": 600.0, "step": 1.0}),
|
|
"max_frames": ("INT", {"default": 60, "min": 1, "max": 512}),
|
|
"overlong_policy": (["trim", "reject"], {"default": "trim"}),
|
|
"include_image_with_video": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"image": ("IMAGE",),
|
|
"video": ("VIDEO",),
|
|
"visual_description": ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True}),
|
|
"media_synthesis_mode": (MEDIA_SYNTHESIS_MODE_CHOICES, {"default": "video_recreation"}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (CONTEXT_TYPE, "STRING", "STRING")
|
|
RETURN_NAMES = ("gemma_context", "context_json", "context_preview")
|
|
FUNCTION = "build"
|
|
CATEGORY = CATEGORY
|
|
|
|
def build(
|
|
self,
|
|
user_prompt: str,
|
|
sample_fps: float = 1.0,
|
|
max_duration_seconds: float = 60.0,
|
|
max_frames: int = 60,
|
|
overlong_policy: str = "trim",
|
|
include_image_with_video: bool = True,
|
|
image=None,
|
|
video=None,
|
|
visual_description: str = "",
|
|
media_synthesis_mode: str = "video_recreation",
|
|
):
|
|
media_context, _metadata_json = DiffusionGemmaMediaSampler().sample(
|
|
sample_fps,
|
|
max_duration_seconds,
|
|
max_frames,
|
|
overlong_policy,
|
|
include_image_with_video,
|
|
image,
|
|
video,
|
|
visual_description,
|
|
media_synthesis_mode,
|
|
)
|
|
context = _gemma_context_from_media(user_prompt, media_context, visual_description)
|
|
if context.images is not None:
|
|
note = "Media pixels are sampled and carried in DG_CONTEXT; the generator decides whether the active backend can consume them."
|
|
if note not in context.warnings:
|
|
context.warnings.append(note)
|
|
context.media_metadata["warnings"] = context.warnings
|
|
return (context, _json_dumps(_gemma_context_to_jsonable(context)), _gemma_context_preview(context))
|
|
|
|
|
|
class DiffusionGemmaReferencePrep:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"image": ("IMAGE",),
|
|
"target_aspect_ratio": (REFERENCE_PREP_ASPECT_CHOICES, {"default": "source"}),
|
|
"aspect_policy": (
|
|
["preserve_source", "center_crop_to_target"],
|
|
{
|
|
"default": "preserve_source",
|
|
"tooltip": "preserve_source keeps the original framing and only scales by long edge. center_crop_to_target crops to the selected aspect ratio before resizing.",
|
|
},
|
|
),
|
|
"long_edge": (
|
|
"INT",
|
|
{
|
|
"default": 1280,
|
|
"min": 256,
|
|
"max": 4096,
|
|
"step": 8,
|
|
"tooltip": "Recommended practical range for DiffusionGemma image ingestion is about 1024-1536 on the long edge.",
|
|
},
|
|
),
|
|
"multiple": (
|
|
"INT",
|
|
{
|
|
"default": 8,
|
|
"min": 1,
|
|
"max": 128,
|
|
"step": 1,
|
|
"tooltip": "Rounds output width and height to this multiple.",
|
|
},
|
|
),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE", "STRING", "INT", "INT", "STRING")
|
|
RETURN_NAMES = ("image", "prep_metadata_json", "width", "height", "resolution_selector_preset")
|
|
FUNCTION = "prepare"
|
|
CATEGORY = OPTIONAL_CATEGORY
|
|
|
|
def prepare(
|
|
self,
|
|
image: Any,
|
|
target_aspect_ratio: str = "source",
|
|
aspect_policy: str = "preserve_source",
|
|
long_edge: int = 1280,
|
|
multiple: int = 8,
|
|
):
|
|
try:
|
|
import torch
|
|
|
|
tensor = image if isinstance(image, torch.Tensor) else torch.as_tensor(image)
|
|
except Exception:
|
|
tensor = image
|
|
|
|
if not hasattr(tensor, "shape") or len(tensor.shape) < 4:
|
|
raise ValueError("DiffusionGemma Reference Prep requires an IMAGE tensor with shape [batch, height, width, channels].")
|
|
|
|
source_height = int(tensor.shape[1])
|
|
source_width = int(tensor.shape[2])
|
|
chosen_policy = str(aspect_policy or "preserve_source")
|
|
chosen_ratio = str(target_aspect_ratio or "source")
|
|
|
|
cropped = tensor
|
|
crop_applied = False
|
|
effective_ratio = chosen_ratio
|
|
|
|
if chosen_ratio != "source" and chosen_policy == "center_crop_to_target":
|
|
crop_width, crop_height = _center_crop_dimensions_for_aspect(source_width, source_height, chosen_ratio)
|
|
if crop_width != source_width or crop_height != source_height:
|
|
cropped = _center_crop_image_batch(tensor, crop_width, crop_height)
|
|
crop_applied = True
|
|
working_width = crop_width
|
|
working_height = crop_height
|
|
target_width, target_height = _aspect_dimensions_for_long_edge(chosen_ratio, long_edge, multiple)
|
|
else:
|
|
working_width = source_width
|
|
working_height = source_height
|
|
effective_ratio = chosen_ratio if chosen_ratio != "source" else _normalize_aspect_ratio_key(f"{source_width}:{source_height}")
|
|
target_width, target_height = _dimensions_for_long_edge(source_width, source_height, long_edge, multiple)
|
|
|
|
prepared = _resize_image_batch(cropped, target_width, target_height)
|
|
|
|
metadata = {
|
|
"source_width": source_width,
|
|
"source_height": source_height,
|
|
"target_aspect_ratio": chosen_ratio,
|
|
"aspect_policy": chosen_policy,
|
|
"crop_applied": crop_applied,
|
|
"working_width": int(working_width),
|
|
"working_height": int(working_height),
|
|
"output_width": int(target_width),
|
|
"output_height": int(target_height),
|
|
"long_edge": int(long_edge),
|
|
"multiple": int(multiple),
|
|
"resolution_selector_preset": _resolution_selector_preset(chosen_ratio)
|
|
if chosen_ratio != "source"
|
|
else _known_aspect_ratio_label_from_dimensions(target_width, target_height),
|
|
"note": "Pre-resizing for DiffusionGemma consistency/efficiency. Preserve-source is safest; center-crop can standardize framing when you intentionally want one aspect ratio.",
|
|
}
|
|
return (
|
|
prepared,
|
|
_json_dumps(metadata),
|
|
int(target_width),
|
|
int(target_height),
|
|
str(metadata["resolution_selector_preset"]),
|
|
)
|
|
|
|
|
|
class DiffusionGemmaTargetProfile:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"target_profile": (["ltx", "ideogram4"], {"default": "ltx"}),
|
|
"audio_mode": (
|
|
["auto_scene_audio", "explicit_sound_design", "visual_only"],
|
|
{"default": "auto_scene_audio"},
|
|
),
|
|
"audio_guidance": ("STRING", {"default": "", "multiline": True}),
|
|
"target_duration_seconds": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 600.0, "step": 0.1}),
|
|
"ltx_style": ("STRING", {"default": "", "multiline": False}),
|
|
"ideogram_aspect_ratio": (
|
|
["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"],
|
|
{"default": "1:1"},
|
|
),
|
|
"ideogram_render_style": ("STRING", {"default": "", "multiline": False}),
|
|
"ideogram_exact_text": ("STRING", {"default": "", "multiline": True}),
|
|
"ideogram_json_output": ("BOOLEAN", {"default": True}),
|
|
"negative_prompt_mode": (["auto", "empty", "custom"], {"default": "auto"}),
|
|
"negative_prompt_guidance": ("STRING", {"default": "", "multiline": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (TARGET_PROFILE_TYPE, "STRING")
|
|
RETURN_NAMES = ("target_profile_config", "target_profile_json")
|
|
FUNCTION = "build"
|
|
CATEGORY = CATEGORY
|
|
|
|
def build(
|
|
self,
|
|
target_profile: str,
|
|
audio_mode: str,
|
|
audio_guidance: str,
|
|
target_duration_seconds: float,
|
|
ltx_style: str,
|
|
ideogram_aspect_ratio: str,
|
|
ideogram_render_style: str,
|
|
ideogram_exact_text: str,
|
|
ideogram_json_output: bool,
|
|
negative_prompt_mode: str,
|
|
negative_prompt_guidance: str,
|
|
):
|
|
config = _make_target_profile_config(
|
|
target_profile,
|
|
audio_mode,
|
|
audio_guidance,
|
|
target_duration_seconds,
|
|
ltx_style,
|
|
ideogram_aspect_ratio,
|
|
ideogram_render_style,
|
|
ideogram_exact_text,
|
|
ideogram_json_output,
|
|
negative_prompt_mode,
|
|
negative_prompt_guidance,
|
|
)
|
|
return (config, _json_dumps(_target_profile_config_to_dict(config)))
|
|
|
|
|
|
class DiffusionGemmaCoTGenerator:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_config": (MODEL_TYPE,),
|
|
"gemma_context": (CONTEXT_TYPE,),
|
|
"target_profile_config": (TARGET_PROFILE_TYPE,),
|
|
"temperature": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.45,
|
|
"min": 0.0,
|
|
"max": 2.0,
|
|
"step": 0.05,
|
|
"tooltip": "Generation randomness for this CoT pass. Rewrites GGUF --temp without changing the loader node.",
|
|
},
|
|
),
|
|
"creativity_mode": (
|
|
["faithful", "editorial", "cinematic", "concept_art", "wild"],
|
|
{"default": "editorial"},
|
|
),
|
|
"creative_strength": ("FLOAT", {"default": 0.6, "min": 0.0, "max": 1.5, "step": 0.05}),
|
|
"thinking_mode": (["auto", "on", "off"], {"default": "auto"}),
|
|
"max_new_tokens": (
|
|
"INT",
|
|
{
|
|
"default": 640,
|
|
"min": 128,
|
|
"max": 2048,
|
|
"step": 64,
|
|
"tooltip": "Output budget. Lower values are faster; higher values help Ideogram JSON or dual-output prompts finish cleanly.",
|
|
},
|
|
),
|
|
},
|
|
"hidden": {
|
|
"unique_id": "UNIQUE_ID",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("final_json", "reasoning_text", "raw_response", "metadata_json")
|
|
FUNCTION = "generate"
|
|
CATEGORY = CATEGORY
|
|
|
|
def generate(
|
|
self,
|
|
model_config: RuntimeConfig,
|
|
gemma_context: GemmaContext,
|
|
target_profile_config: TargetProfileConfig,
|
|
temperature: float,
|
|
creativity_mode: str,
|
|
creative_strength: float,
|
|
thinking_mode: str,
|
|
max_new_tokens: int = 640,
|
|
unique_id: str | None = None,
|
|
):
|
|
runtime_config = _runtime_config_with_temperature(model_config, temperature)
|
|
packet, raw_output, reasoning_text, _fallback_reason, _parse_valid, _salvage_warning, metadata = _run_generation_packet(
|
|
runtime_config,
|
|
gemma_context,
|
|
target_profile_config,
|
|
DEFAULT_MASTER_PROMPT,
|
|
True,
|
|
int(max_new_tokens),
|
|
8000,
|
|
creativity_mode,
|
|
creative_strength,
|
|
thinking_mode,
|
|
node_id=str(unique_id) if unique_id is not None else None,
|
|
)
|
|
return (_json_dumps(packet), reasoning_text, raw_output, _json_dumps(metadata))
|
|
|
|
|
|
class DiffusionGemmaJSONSplitter:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"final_json": ("STRING", {"default": "", "multiline": True, "forceInput": True}),
|
|
},
|
|
"optional": {
|
|
"gemma_context": (CONTEXT_TYPE,),
|
|
"target_profile_config": (TARGET_PROFILE_TYPE,),
|
|
"fallback_prompt": ("STRING", {"default": "", "multiline": True}),
|
|
"resolution_megapixels": (
|
|
"FLOAT",
|
|
{
|
|
"default": 2.0,
|
|
"min": 0.1,
|
|
"max": 16.0,
|
|
"step": 0.1,
|
|
"tooltip": "Used only for the splitter's resolution_width/resolution_height outputs. Matches Resolution Selector megapixels.",
|
|
},
|
|
),
|
|
"resolution_multiple": (
|
|
"INT",
|
|
{
|
|
"default": 8,
|
|
"min": 8,
|
|
"max": 128,
|
|
"step": 4,
|
|
"tooltip": "Used only for the splitter's resolution_width/resolution_height outputs. Matches Resolution Selector multiple.",
|
|
},
|
|
),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = (
|
|
"STRING",
|
|
"STRING",
|
|
"STRING",
|
|
"STRING",
|
|
"STRING",
|
|
"STRING",
|
|
"STRING",
|
|
"STRING",
|
|
"BOOLEAN",
|
|
"STRING",
|
|
"INT",
|
|
"INT",
|
|
"BOOLEAN",
|
|
)
|
|
RETURN_NAMES = (
|
|
"ltx_prompt",
|
|
"ideogram_prompt",
|
|
"negative_prompt",
|
|
"aspect_ratio",
|
|
"metadata_json",
|
|
"scene_segments_json",
|
|
"local_prompts",
|
|
"segment_lengths",
|
|
"is_valid",
|
|
"resolution_selector_preset",
|
|
"resolution_width",
|
|
"resolution_height",
|
|
"ready_for_generation",
|
|
)
|
|
FUNCTION = "split"
|
|
CATEGORY = CATEGORY
|
|
|
|
def split(
|
|
self,
|
|
final_json: str,
|
|
gemma_context: GemmaContext | None = None,
|
|
target_profile_config: TargetProfileConfig | None = None,
|
|
fallback_prompt: str = "",
|
|
resolution_megapixels: float = 2.0,
|
|
resolution_multiple: int = 8,
|
|
):
|
|
split_start = time.perf_counter()
|
|
context = gemma_context if isinstance(gemma_context, GemmaContext) else _gemma_context_from_media(fallback_prompt)
|
|
target = target_profile_config if isinstance(target_profile_config, TargetProfileConfig) else _make_target_profile_config()
|
|
cleaned = _strip_markdown(_strip_thinking(str(final_json or "")))
|
|
strict_packet: dict[str, Any] | None = None
|
|
try:
|
|
parsed = json.loads(cleaned)
|
|
if isinstance(parsed, dict):
|
|
strict_packet = parsed
|
|
except Exception:
|
|
strict_packet = None
|
|
if strict_packet is None:
|
|
packet, is_valid, warning = repair_or_salvage_prompt_packet(final_json, context, target, 8000)
|
|
else:
|
|
packet, is_valid, warning = strict_packet, True, ""
|
|
packet_metadata_for_split = packet.setdefault("metadata", {}) if isinstance(packet, dict) else {}
|
|
if isinstance(packet_metadata_for_split, dict):
|
|
packet_metadata_for_split["json_parse_valid"] = bool(strict_packet is not None)
|
|
if warning:
|
|
packet_metadata_for_split["salvage_warning"] = warning
|
|
ltx_prompt, ideogram_prompt, negative_prompt, scene_segments_json, local_prompts, segment_lengths, aspect_ratio, packet_metadata = _packet_to_prompt_outputs(
|
|
packet,
|
|
context,
|
|
target,
|
|
8000,
|
|
)
|
|
resolution_selector_preset = _resolution_selector_preset(aspect_ratio)
|
|
resolution_width, resolution_height = _resolution_selector_dimensions(
|
|
aspect_ratio,
|
|
resolution_megapixels,
|
|
resolution_multiple,
|
|
)
|
|
metadata = _extract_json_object(packet_metadata) or {}
|
|
ready_for_generation = bool(metadata.get("ready_for_generation"))
|
|
blocked_reasons = metadata.get("blocked_reasons") if isinstance(metadata.get("blocked_reasons"), list) else []
|
|
if not ready_for_generation:
|
|
ltx_prompt = ""
|
|
ideogram_prompt = ""
|
|
negative_prompt = ""
|
|
scene_segments_json = "[]"
|
|
local_prompts = ""
|
|
segment_lengths = ""
|
|
metadata.update(
|
|
{
|
|
"splitter_json_valid": bool(is_valid and ready_for_generation),
|
|
"splitter_warning": warning,
|
|
"target_profile": _target_profile_config_to_dict(target)["target_profile"],
|
|
"aspect_ratio": aspect_ratio,
|
|
"ready_for_generation": ready_for_generation,
|
|
"blocked_reasons": blocked_reasons,
|
|
"resolution_selector_preset": resolution_selector_preset,
|
|
"resolution_width": resolution_width,
|
|
"resolution_height": resolution_height,
|
|
"resolution_megapixels": max(0.1, min(16.0, _safe_float(resolution_megapixels, 2.0))),
|
|
"resolution_multiple": int(max(8, min(128, _safe_float(resolution_multiple, 8)))),
|
|
"splitter_timing": {
|
|
"split_seconds": round(time.perf_counter() - split_start, 3),
|
|
"input_chars": len(str(final_json or "")),
|
|
},
|
|
}
|
|
)
|
|
return (
|
|
ltx_prompt,
|
|
ideogram_prompt,
|
|
negative_prompt,
|
|
aspect_ratio,
|
|
_json_dumps(metadata),
|
|
scene_segments_json,
|
|
local_prompts,
|
|
segment_lengths,
|
|
bool(is_valid and ready_for_generation),
|
|
resolution_selector_preset,
|
|
resolution_width,
|
|
resolution_height,
|
|
ready_for_generation,
|
|
)
|
|
|
|
|
|
class DiffusionGemmaPromptBuilder:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_config": (MODEL_TYPE,),
|
|
"user_prompt": ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True}),
|
|
"master_prompt": ("STRING", {"default": DEFAULT_MASTER_PROMPT, "multiline": True}),
|
|
"target_profile": (["ltx", "ideogram4"], {"default": "ltx"}),
|
|
"runtime_required": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": False,
|
|
"tooltip": "If true, fail instead of using deterministic template output when DiffusionGemma is unavailable.",
|
|
},
|
|
),
|
|
"max_new_tokens": ("INT", {"default": 768, "min": 64, "max": 4096}),
|
|
"max_output_chars": ("INT", {"default": 4000, "min": 512, "max": 20000}),
|
|
"audio_mode": (
|
|
["auto_scene_audio", "explicit_sound_design", "visual_only"],
|
|
{
|
|
"default": "auto_scene_audio",
|
|
"tooltip": "Controls how DiffusionGemma writes audio into the LTX prompt before generation.",
|
|
},
|
|
),
|
|
"audio_guidance": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": True,
|
|
"tooltip": "Optional sound design instruction for DiffusionGemma to weave into ltx_prompt.",
|
|
},
|
|
),
|
|
"target_duration_seconds": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 600.0,
|
|
"step": 0.1,
|
|
"tooltip": "Optional target duration for the direct LTX prompt. Video media duration wins only when this is 0.",
|
|
},
|
|
),
|
|
"ltx_style": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": False,
|
|
"tooltip": "Optional LTX style phrase for the Prompt Builder to weave into ltx_prompt.",
|
|
},
|
|
),
|
|
"ideogram_aspect_ratio": (
|
|
["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"],
|
|
{"default": "1:1"},
|
|
),
|
|
"ideogram_render_style": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": False,
|
|
"tooltip": "Optional Ideogram4 visual style or medium.",
|
|
},
|
|
),
|
|
"ideogram_exact_text": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": True,
|
|
"tooltip": "Optional exact text that Ideogram4 should render verbatim.",
|
|
},
|
|
),
|
|
"ideogram_json_output": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
"tooltip": "If true, ideogram_prompt is returned as structured Ideogram4 JSON text.",
|
|
},
|
|
),
|
|
"creativity_mode": (
|
|
["faithful", "editorial", "cinematic", "concept_art", "wild"],
|
|
{
|
|
"default": "editorial",
|
|
"tooltip": "Controls how boldly DiffusionGemma adds art direction while preserving the requested subject and exact text.",
|
|
},
|
|
),
|
|
"creative_strength": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.6,
|
|
"min": 0.0,
|
|
"max": 1.5,
|
|
"step": 0.05,
|
|
"tooltip": "Amount of added detail. 0 is faithful; 0.6 editorial; 1.0+ more concept-art/wild.",
|
|
},
|
|
),
|
|
"negative_prompt_mode": (
|
|
["auto", "empty", "custom"],
|
|
{
|
|
"default": "auto",
|
|
"tooltip": "auto lets the model emit useful negatives; empty forces none; custom returns your supplied negative prompt exactly.",
|
|
},
|
|
),
|
|
"negative_prompt_guidance": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": True,
|
|
"tooltip": "Used only when negative_prompt_mode is custom, or as guidance for auto negatives.",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"media_context": (MEDIA_TYPE,),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING", "STRING", "STRING", "STRING")
|
|
RETURN_NAMES = (
|
|
"ltx_prompt",
|
|
"ideogram_prompt",
|
|
"negative_prompt",
|
|
"metadata_json",
|
|
"scene_segments_json",
|
|
"local_prompts",
|
|
"segment_lengths",
|
|
)
|
|
FUNCTION = "build"
|
|
CATEGORY = COMPATIBILITY_CATEGORY
|
|
|
|
def build(
|
|
self,
|
|
model_config: RuntimeConfig,
|
|
user_prompt: str,
|
|
master_prompt: str,
|
|
target_profile: str,
|
|
runtime_required: bool,
|
|
max_new_tokens: int,
|
|
max_output_chars: int,
|
|
audio_mode: str = "auto_scene_audio",
|
|
audio_guidance: str = "",
|
|
target_duration_seconds: float = 0.0,
|
|
ltx_style: str = "",
|
|
ideogram_aspect_ratio: str = "1:1",
|
|
ideogram_render_style: str = "",
|
|
ideogram_exact_text: str = "",
|
|
ideogram_json_output: bool = True,
|
|
creativity_mode: str = "editorial",
|
|
creative_strength: float = 0.6,
|
|
negative_prompt_mode: str = "auto",
|
|
negative_prompt_guidance: str = "",
|
|
media_context: MediaContext | None = None,
|
|
):
|
|
if not isinstance(model_config, RuntimeConfig):
|
|
raise TypeError("model_config must come from DiffusionGemma Model Loader.")
|
|
|
|
context = _gemma_context_from_media(user_prompt, media_context)
|
|
target_config_obj = _make_target_profile_config(
|
|
target_profile,
|
|
audio_mode,
|
|
audio_guidance,
|
|
target_duration_seconds,
|
|
ltx_style,
|
|
ideogram_aspect_ratio,
|
|
ideogram_render_style,
|
|
ideogram_exact_text,
|
|
ideogram_json_output,
|
|
negative_prompt_mode,
|
|
negative_prompt_guidance,
|
|
)
|
|
packet, raw_output, reasoning_text, fallback_reason, parse_valid, salvage_warning, metadata = _run_generation_packet(
|
|
model_config,
|
|
context,
|
|
target_config_obj,
|
|
master_prompt,
|
|
bool(runtime_required),
|
|
int(max_new_tokens),
|
|
int(max_output_chars),
|
|
creativity_mode,
|
|
creative_strength,
|
|
"off",
|
|
)
|
|
(
|
|
ltx_prompt,
|
|
ideogram_prompt,
|
|
negative_prompt,
|
|
scene_segments_json,
|
|
local_prompts,
|
|
segment_lengths,
|
|
_aspect_ratio,
|
|
_packet_metadata,
|
|
) = _packet_to_prompt_outputs(
|
|
packet,
|
|
context,
|
|
target_config_obj,
|
|
int(max_output_chars),
|
|
_normalize_creativity_mode(creativity_mode),
|
|
_normalize_creative_strength(creative_strength),
|
|
)
|
|
metadata.update(
|
|
{
|
|
"raw_output_available": bool(raw_output),
|
|
"reasoning_available": bool(reasoning_text),
|
|
"used_template_fallback": bool(fallback_reason),
|
|
"json_parse_valid": bool(parse_valid),
|
|
"salvage_warning": salvage_warning,
|
|
}
|
|
)
|
|
if not bool(metadata.get("ready_for_generation")):
|
|
ltx_prompt = ""
|
|
ideogram_prompt = ""
|
|
negative_prompt = ""
|
|
scene_segments_json = "[]"
|
|
local_prompts = ""
|
|
segment_lengths = ""
|
|
return (
|
|
ltx_prompt,
|
|
ideogram_prompt,
|
|
negative_prompt,
|
|
_json_dumps(metadata),
|
|
scene_segments_json,
|
|
local_prompts,
|
|
segment_lengths,
|
|
)
|
|
|
|
target_profile = _normalize_target_profile(target_profile)
|
|
media_metadata = _metadata_from_media(media_context)
|
|
duration = _safe_float(media_metadata.get("duration_seconds"), 0.0)
|
|
audio_mode = _normalize_audio_mode(audio_mode)
|
|
audio_guidance = _audio_guidance_text(audio_guidance)
|
|
creativity_mode = _normalize_creativity_mode(creativity_mode)
|
|
creative_strength = _normalize_creative_strength(creative_strength)
|
|
negative_prompt_mode = _normalize_negative_prompt_mode(negative_prompt_mode)
|
|
negative_prompt_guidance = _negative_prompt_guidance_text(negative_prompt_guidance)
|
|
effective_duration = _effective_duration_seconds(media_metadata, target_duration_seconds)
|
|
model_prompt = _build_model_prompt(
|
|
user_prompt,
|
|
master_prompt,
|
|
target_profile,
|
|
media_metadata,
|
|
audio_mode,
|
|
audio_guidance,
|
|
target_duration_seconds,
|
|
ltx_style,
|
|
ideogram_aspect_ratio,
|
|
ideogram_render_style,
|
|
ideogram_exact_text,
|
|
ideogram_json_output,
|
|
creativity_mode,
|
|
creative_strength,
|
|
negative_prompt_mode,
|
|
negative_prompt_guidance,
|
|
)
|
|
raw_output = ""
|
|
fallback_reason = ""
|
|
packet: dict[str, Any]
|
|
|
|
try:
|
|
if model_config.backend == "template":
|
|
raise RuntimeError("Template backend selected.")
|
|
raw_output = _run_backend(model_config, model_prompt, media_context, int(max_new_tokens))
|
|
parsed = _extract_json_object(raw_output)
|
|
if parsed is None:
|
|
salvaged = _packet_from_plain_text_output(
|
|
raw_output,
|
|
user_prompt,
|
|
target_profile,
|
|
media_metadata,
|
|
duration,
|
|
audio_mode,
|
|
audio_guidance,
|
|
target_duration_seconds,
|
|
ltx_style,
|
|
ideogram_aspect_ratio,
|
|
ideogram_render_style,
|
|
ideogram_exact_text,
|
|
ideogram_json_output,
|
|
creativity_mode,
|
|
creative_strength,
|
|
negative_prompt_mode,
|
|
negative_prompt_guidance,
|
|
int(max_output_chars),
|
|
)
|
|
if salvaged is None:
|
|
raise RuntimeError("Model output did not contain valid JSON or usable plain text.")
|
|
packet = salvaged
|
|
else:
|
|
packet = parsed
|
|
except Exception as exc:
|
|
fallback_reason = str(exc)
|
|
if model_config.backend == "transformers_inprocess":
|
|
_release_transformers_runtime()
|
|
if runtime_required:
|
|
raise RuntimeError(f"DiffusionGemma runtime failed and runtime_required is true: {exc}") from exc
|
|
packet = _template_packet(
|
|
user_prompt,
|
|
target_profile,
|
|
media_metadata,
|
|
duration,
|
|
audio_mode,
|
|
audio_guidance,
|
|
target_duration_seconds,
|
|
ltx_style,
|
|
ideogram_aspect_ratio,
|
|
ideogram_render_style,
|
|
ideogram_exact_text,
|
|
ideogram_json_output,
|
|
creativity_mode,
|
|
creative_strength,
|
|
negative_prompt_mode,
|
|
negative_prompt_guidance,
|
|
)
|
|
packet.setdefault("metadata", {})["fallback_reason"] = fallback_reason
|
|
finally:
|
|
_maybe_unload(model_config)
|
|
|
|
prompt_fallback = _clean_prompt_request_text(user_prompt)
|
|
ltx_prompt = _prompt_field_to_text(packet.get("ltx_prompt", ""), prompt_fallback, int(max_output_chars))
|
|
ideogram_value = packet.get("ideogram_prompt", "")
|
|
if not ideogram_value and {"high_level_description", "compositional_deconstruction"}.issubset(packet.keys()):
|
|
ideogram_value = packet
|
|
ideogram_prompt = _ideogram_prompt_value_to_text(
|
|
ideogram_value,
|
|
prompt_fallback,
|
|
ideogram_aspect_ratio,
|
|
ideogram_render_style,
|
|
ideogram_exact_text,
|
|
bool(ideogram_json_output),
|
|
int(max_output_chars),
|
|
creativity_mode,
|
|
creative_strength,
|
|
)
|
|
negative_prompt = _apply_negative_prompt_policy(
|
|
packet.get("negative_prompt", ""),
|
|
negative_prompt_mode,
|
|
negative_prompt_guidance,
|
|
)
|
|
scene_segments = packet.get("scene_segments")
|
|
if not isinstance(scene_segments, list) or not scene_segments:
|
|
scene_segments = _split_segments(ltx_prompt, effective_duration or duration)
|
|
local_prompts, segment_lengths = _segments_to_director_strings(scene_segments)
|
|
metadata = packet.get("metadata") if isinstance(packet.get("metadata"), dict) else {}
|
|
metadata.update(
|
|
{
|
|
"backend": model_config.backend,
|
|
"runtime_ready": bool(model_config.status.get("ready")),
|
|
"raw_output_available": bool(raw_output),
|
|
"used_template_fallback": bool(fallback_reason),
|
|
"media": media_metadata,
|
|
"audio_mode": audio_mode,
|
|
"audio_guidance": audio_guidance,
|
|
"target_duration_seconds": effective_duration,
|
|
"ltx_style": _clean_control_text(ltx_style, 240),
|
|
"ideogram_aspect_ratio": ideogram_aspect_ratio,
|
|
"ideogram_render_style": _clean_control_text(ideogram_render_style, 240),
|
|
"ideogram_exact_text": _sanitize_prompt_text(
|
|
ideogram_exact_text,
|
|
"",
|
|
1000,
|
|
strip_thinking=True,
|
|
strip_markdown=True,
|
|
),
|
|
"ideogram_json_output": bool(ideogram_json_output),
|
|
"creativity_mode": creativity_mode,
|
|
"creative_strength": creative_strength,
|
|
"negative_prompt_mode": negative_prompt_mode,
|
|
"negative_prompt_guidance": negative_prompt_guidance,
|
|
}
|
|
)
|
|
return (
|
|
ltx_prompt,
|
|
ideogram_prompt,
|
|
negative_prompt,
|
|
_json_dumps(metadata),
|
|
_json_dumps(scene_segments),
|
|
local_prompts,
|
|
segment_lengths,
|
|
)
|
|
|
|
|
|
class DiffusionGemmaSimplePromptBuilder:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"model_config": (MODEL_TYPE,),
|
|
"user_prompt": ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True}),
|
|
"target_profile": (["ltx", "ideogram4"], {"default": "ltx"}),
|
|
"audio_mode": (
|
|
["auto_scene_audio", "explicit_sound_design", "visual_only"],
|
|
{
|
|
"default": "auto_scene_audio",
|
|
"tooltip": "Controls whether LTX prompts include woven-in sound design.",
|
|
},
|
|
),
|
|
"audio_guidance": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": True,
|
|
"tooltip": "Optional sound design instruction for DiffusionGemma to weave into ltx_prompt.",
|
|
},
|
|
),
|
|
"target_duration_seconds": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.0,
|
|
"min": 0.0,
|
|
"max": 600.0,
|
|
"step": 0.1,
|
|
"tooltip": "Optional target duration for the direct LTX prompt. Video media duration wins only when this is 0.",
|
|
},
|
|
),
|
|
"ltx_style": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": False,
|
|
"tooltip": "Optional LTX style phrase for the Prompt Builder to weave into ltx_prompt.",
|
|
},
|
|
),
|
|
"ideogram_aspect_ratio": (
|
|
["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"],
|
|
{"default": "1:1"},
|
|
),
|
|
"ideogram_render_style": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": False,
|
|
"tooltip": "Optional Ideogram4 visual style or medium.",
|
|
},
|
|
),
|
|
"ideogram_exact_text": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": True,
|
|
"tooltip": "Optional exact text that Ideogram4 should render verbatim.",
|
|
},
|
|
),
|
|
"ideogram_json_output": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
"tooltip": "If true, ideogram_prompt is returned as structured Ideogram4 JSON text.",
|
|
},
|
|
),
|
|
"creativity_mode": (
|
|
["faithful", "editorial", "cinematic", "concept_art", "wild"],
|
|
{
|
|
"default": "editorial",
|
|
"tooltip": "Controls how boldly DiffusionGemma adds art direction while preserving the requested subject and exact text.",
|
|
},
|
|
),
|
|
"creative_strength": (
|
|
"FLOAT",
|
|
{
|
|
"default": 0.6,
|
|
"min": 0.0,
|
|
"max": 1.5,
|
|
"step": 0.05,
|
|
"tooltip": "Amount of added detail. 0 is faithful; 0.6 editorial; 1.0+ more concept-art/wild.",
|
|
},
|
|
),
|
|
"negative_prompt_mode": (
|
|
["auto", "empty", "custom"],
|
|
{
|
|
"default": "auto",
|
|
"tooltip": "auto lets the model emit useful negatives; empty forces none; custom returns your supplied negative prompt exactly.",
|
|
},
|
|
),
|
|
"negative_prompt_guidance": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": True,
|
|
"tooltip": "Used only when negative_prompt_mode is custom, or as guidance for auto negatives.",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"media_context": (MEDIA_TYPE,),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = DiffusionGemmaPromptBuilder.RETURN_TYPES
|
|
RETURN_NAMES = DiffusionGemmaPromptBuilder.RETURN_NAMES
|
|
FUNCTION = "build"
|
|
CATEGORY = COMPATIBILITY_CATEGORY
|
|
|
|
def build(
|
|
self,
|
|
model_config: RuntimeConfig,
|
|
user_prompt: str,
|
|
target_profile: str,
|
|
audio_mode: str = "auto_scene_audio",
|
|
audio_guidance: str = "",
|
|
target_duration_seconds: float = 0.0,
|
|
ltx_style: str = "",
|
|
ideogram_aspect_ratio: str = "1:1",
|
|
ideogram_render_style: str = "",
|
|
ideogram_exact_text: str = "",
|
|
ideogram_json_output: bool = True,
|
|
creativity_mode: str = "editorial",
|
|
creative_strength: float = 0.6,
|
|
negative_prompt_mode: str = "auto",
|
|
negative_prompt_guidance: str = "",
|
|
media_context: MediaContext | None = None,
|
|
):
|
|
return DiffusionGemmaPromptBuilder().build(
|
|
model_config,
|
|
user_prompt,
|
|
DEFAULT_MASTER_PROMPT,
|
|
target_profile,
|
|
False,
|
|
768,
|
|
8000,
|
|
audio_mode,
|
|
audio_guidance,
|
|
target_duration_seconds,
|
|
ltx_style,
|
|
ideogram_aspect_ratio,
|
|
ideogram_render_style,
|
|
ideogram_exact_text,
|
|
ideogram_json_output,
|
|
creativity_mode,
|
|
creative_strength,
|
|
negative_prompt_mode,
|
|
negative_prompt_guidance,
|
|
media_context,
|
|
)
|
|
|
|
|
|
class LTX2PromptAdapter:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"base_prompt": ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True}),
|
|
"duration_seconds": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 600.0, "step": 0.1}),
|
|
"style": ("STRING", {"default": "", "multiline": False}),
|
|
"include_audio": (
|
|
"BOOLEAN",
|
|
{
|
|
"default": True,
|
|
"tooltip": "If true, the adapter may add a final scene-specific audio cue when one is missing. If false, it leaves audio untouched; use Prompt Builder audio_mode=visual_only to request no audio upstream.",
|
|
},
|
|
),
|
|
"max_segments": ("INT", {"default": 6, "min": 1, "max": 24}),
|
|
"audio_guidance": (
|
|
"STRING",
|
|
{
|
|
"default": "",
|
|
"multiline": True,
|
|
"tooltip": "Optional final LTX sound design cue. For best results, set audio on DiffusionGemma Prompt Builder so the model writes it upstream.",
|
|
},
|
|
),
|
|
},
|
|
"optional": {
|
|
"metadata_json": ("STRING", {"default": "", "multiline": True}),
|
|
"scene_segments_json": ("STRING", {"default": "", "multiline": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING", "STRING", "STRING")
|
|
RETURN_NAMES = ("ltx_prompt", "local_prompts", "segment_lengths", "metadata_json")
|
|
FUNCTION = "adapt"
|
|
CATEGORY = OPTIONAL_CATEGORY
|
|
|
|
def adapt(
|
|
self,
|
|
base_prompt: str,
|
|
duration_seconds: float,
|
|
style: str,
|
|
include_audio: bool,
|
|
max_segments: int,
|
|
audio_guidance: str = "",
|
|
metadata_json: str = "",
|
|
scene_segments_json: str = "",
|
|
):
|
|
metadata = _extract_json_object(metadata_json) or {}
|
|
media = metadata.get("media") if isinstance(metadata.get("media"), dict) else {}
|
|
effective_duration = float(duration_seconds) if float(duration_seconds) > 0 else _safe_float(media.get("duration_seconds"), 0.0)
|
|
audio_guidance = _audio_guidance_text(audio_guidance)
|
|
source_prompt = _single_paragraph(base_prompt)
|
|
source_segments = _load_scene_segments(scene_segments_json)
|
|
if not source_prompt and source_segments:
|
|
source_prompt = _single_paragraph(
|
|
" ".join(
|
|
str(segment.get("prompt", "")) if isinstance(segment, dict) else str(segment)
|
|
for segment in source_segments
|
|
)
|
|
)
|
|
prompt = _apply_ltx2_controls(source_prompt, effective_duration, style, bool(include_audio), audio_guidance)
|
|
if not prompt:
|
|
metadata.update(
|
|
{
|
|
"adapter": "ltx",
|
|
"duration_seconds": effective_duration,
|
|
"include_audio": bool(include_audio),
|
|
"audio_guidance": audio_guidance,
|
|
"max_segments": int(max_segments),
|
|
"segment_count": 0,
|
|
"style": _sanitize_prompt_text(style, "", 300, strip_thinking=True, strip_markdown=True).rstrip("."),
|
|
"warning": "LTX adapter received no base_prompt or scene_segments_json, so it returned an empty prompt instead of inventing one.",
|
|
}
|
|
)
|
|
return ("", "", "", _json_dumps(metadata))
|
|
|
|
segments = _normalize_segments(source_segments, source_prompt, effective_duration, int(max_segments))
|
|
segment_style = _sanitize_prompt_text(style, "", 300, strip_thinking=True, strip_markdown=True).rstrip(".")
|
|
if segment_style:
|
|
style_prefix = f"Style: {segment_style}."
|
|
for segment in segments:
|
|
segment_prompt = str(segment.get("prompt", ""))
|
|
if not segment_prompt.lower().startswith(style_prefix.lower()):
|
|
segment["prompt"] = _single_paragraph(f"{style_prefix} {segment_prompt}")
|
|
local_prompts, segment_lengths = _segments_to_director_strings(segments)
|
|
metadata.update(
|
|
{
|
|
"adapter": "ltx",
|
|
"duration_seconds": effective_duration,
|
|
"include_audio": bool(include_audio),
|
|
"audio_guidance": audio_guidance,
|
|
"max_segments": int(max_segments),
|
|
"segment_count": len(segments),
|
|
"style": segment_style,
|
|
}
|
|
)
|
|
return (prompt, local_prompts, segment_lengths, _json_dumps(metadata))
|
|
|
|
|
|
class Ideogram4PromptAdapter:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"base_prompt": ("STRING", {"default": "", "multiline": True, "dynamicPrompts": True}),
|
|
"aspect_ratio": (
|
|
["1:1", "16:9", "9:16", "4:3", "3:4", "3:2", "2:3"],
|
|
{"default": "1:1"},
|
|
),
|
|
"render_style": ("STRING", {"default": "", "multiline": False}),
|
|
"exact_text": ("STRING", {"default": "", "multiline": True}),
|
|
"json_style_output": ("BOOLEAN", {"default": True}),
|
|
},
|
|
"optional": {
|
|
"metadata_json": ("STRING", {"default": "", "multiline": True}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("ideogram_prompt", "metadata_json")
|
|
FUNCTION = "adapt"
|
|
CATEGORY = OPTIONAL_CATEGORY
|
|
|
|
def adapt(
|
|
self,
|
|
base_prompt: str,
|
|
aspect_ratio: str,
|
|
render_style: str,
|
|
exact_text: str,
|
|
json_style_output: bool,
|
|
metadata_json: str = "",
|
|
):
|
|
prompt = _single_paragraph(base_prompt)
|
|
metadata = _extract_json_object(metadata_json) or {}
|
|
metadata["adapter"] = "ideogram4"
|
|
metadata["aspect_ratio"] = aspect_ratio
|
|
metadata["exact_text"] = exact_text.strip()
|
|
metadata["json_style_output"] = bool(json_style_output)
|
|
return (
|
|
_ideogram_prompt_value_to_text(
|
|
prompt,
|
|
prompt,
|
|
aspect_ratio,
|
|
render_style,
|
|
exact_text.strip(),
|
|
bool(json_style_output),
|
|
20000,
|
|
),
|
|
_json_dumps(metadata),
|
|
)
|
|
|
|
|
|
class PromptSanitizer:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"prompt": ("STRING", {"default": "", "multiline": True}),
|
|
"fallback_prompt": ("STRING", {"default": "", "multiline": True}),
|
|
"mode": (["plain_prompt", "json", "ltx", "ideogram4"], {"default": "plain_prompt"}),
|
|
"max_chars": ("INT", {"default": 4000, "min": 1, "max": 20000}),
|
|
"strip_thinking": ("BOOLEAN", {"default": True}),
|
|
"strip_markdown": ("BOOLEAN", {"default": True}),
|
|
"require_json": ("BOOLEAN", {"default": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "BOOLEAN", "STRING")
|
|
RETURN_NAMES = ("sanitized_prompt", "is_valid", "metadata_json")
|
|
FUNCTION = "sanitize"
|
|
CATEGORY = OPTIONAL_CATEGORY
|
|
|
|
def sanitize(
|
|
self,
|
|
prompt: str,
|
|
fallback_prompt: str,
|
|
mode: str,
|
|
max_chars: int,
|
|
strip_thinking: bool,
|
|
strip_markdown: bool,
|
|
require_json: bool,
|
|
):
|
|
value = _sanitize_prompt_text(prompt, fallback_prompt, int(max_chars), bool(strip_thinking), bool(strip_markdown))
|
|
valid = True
|
|
metadata: dict[str, Any] = {"mode": mode, "used_fallback": False}
|
|
if require_json or mode == "json":
|
|
parsed = _extract_json_object(value)
|
|
if parsed is None:
|
|
value = _sanitize_prompt_text(fallback_prompt, "", int(max_chars), bool(strip_thinking), bool(strip_markdown))
|
|
valid = False
|
|
metadata["used_fallback"] = True
|
|
metadata["error"] = "Prompt did not contain valid JSON."
|
|
else:
|
|
value = _json_dumps(parsed)
|
|
if mode in {"ltx", "ltx2", "ltx2_3"}:
|
|
value = _single_paragraph(value)
|
|
if mode == "ideogram4":
|
|
value = _single_paragraph(value)
|
|
return (value, valid, _json_dumps(metadata))
|
|
|
|
|
|
class TextOutput:
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"forceInput": True, "multiline": True}),
|
|
},
|
|
"optional": {
|
|
"label": ("STRING", {"default": "DiffusionGemma output", "multiline": False}),
|
|
"save_to_file": ("BOOLEAN", {"default": True}),
|
|
"filename_prefix": ("STRING", {"default": "diffusiongemma_prompt", "multiline": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING", "STRING")
|
|
RETURN_NAMES = ("text", "saved_path")
|
|
FUNCTION = "show"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = OPTIONAL_CATEGORY
|
|
|
|
def show(
|
|
self,
|
|
text: str,
|
|
label: str = "DiffusionGemma output",
|
|
save_to_file: bool = True,
|
|
filename_prefix: str = "diffusiongemma_prompt",
|
|
):
|
|
title = _sanitize_prompt_text(label, "DiffusionGemma output", 160, strip_thinking=False, strip_markdown=True)
|
|
value = str(text or "")
|
|
saved_path = ""
|
|
if save_to_file:
|
|
try:
|
|
import folder_paths
|
|
|
|
output_dir = Path(folder_paths.get_output_directory()) / "diffusiongemma"
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
safe_prefix = re.sub(r"[^A-Za-z0-9_.-]+", "_", filename_prefix.strip() or "diffusiongemma_prompt")
|
|
counter = 1
|
|
while True:
|
|
candidate = output_dir / f"{safe_prefix}_{counter:05}.txt"
|
|
if not candidate.exists():
|
|
candidate.write_text(value, encoding="utf-8")
|
|
saved_path = str(candidate)
|
|
break
|
|
counter += 1
|
|
except Exception as exc:
|
|
saved_path = f"Save failed: {exc}"
|
|
ui_text = f"{title}\n\n{value}".strip()
|
|
if saved_path:
|
|
ui_text = f"{ui_text}\n\nSaved: {saved_path}"
|
|
return {"ui": {"text": [ui_text]}, "result": (value, saved_path)}
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DiffusionGemmaModelLoader": DiffusionGemmaModelLoader,
|
|
"DiffusionGemmaContextHub": DiffusionGemmaContextHub,
|
|
"DiffusionGemmaTargetProfile": DiffusionGemmaTargetProfile,
|
|
"DiffusionGemmaCoTGenerator": DiffusionGemmaCoTGenerator,
|
|
"DiffusionGemmaJSONSplitter": DiffusionGemmaJSONSplitter,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DiffusionGemmaModelLoader": "DiffusionGemma Model Loader (Advanced)",
|
|
"DiffusionGemmaContextHub": "DiffusionGemma Context Hub",
|
|
"DiffusionGemmaTargetProfile": "DiffusionGemma Target Profile",
|
|
"DiffusionGemmaCoTGenerator": "DiffusionGemma CoT Generator",
|
|
"DiffusionGemmaJSONSplitter": "DiffusionGemma JSON Splitter",
|
|
}
|