Files

1335 lines
47 KiB
Python

import os
import sys
import torch
import numpy as np
import random
import subprocess
from typing import Dict, Tuple, Optional, Any, Union
from PIL import Image
from folder_paths import folder_names_and_paths, models_dir as comfy_models_dir
from comfy.utils import ProgressBar
# Add current directory to path
current_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, current_dir)
# Import BAGEL related modules
try:
from accelerate import (
infer_auto_device_map,
load_checkpoint_and_dispatch,
init_empty_weights,
dispatch_model,
)
from accelerate.utils import BnbQuantizationConfig, load_and_quantize_model
from data.data_utils import add_special_tokens, pil_img2rgb
from data.transforms import ImageTransform
from inferencer import InterleaveInferencer
from modeling.autoencoder import load_ae
from modeling.bagel.qwen2_navit import NaiveCache
from modeling.bagel import (
BagelConfig,
Bagel,
Qwen2Config,
Qwen2ForCausalLM,
SiglipVisionConfig,
SiglipVisionModel,
)
from modeling.qwen2 import Qwen2Tokenizer
except ImportError as e:
print(f"Error importing BAGEL modules: {e}")
print("Please ensure BAGEL model files are properly installed.")
try:
from dfloat11 import DFloat11Model
except ImportError:
print("DFloat11Model not found. DFloat11 support will be unavailable.")
print(
"Please install DFloat11 if you intend to use DFloat11 models: pip install dfloat11"
)
DFloat11Model = None
# Register the BAGEL model folder
folder_names_and_paths["bagel"] = (
[os.path.join(comfy_models_dir, "bagel")],
[".json", ".safetensors"],
)
def set_seed(seed: int) -> int:
"""Set random seeds for reproducibility"""
if seed > 0:
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
return seed
def download_model_with_git(
model_dir: str, repo_id: str = "ByteDance-Seed/BAGEL-7B-MoT"
) -> str:
"""
Download model using git lfs (recommended method)
Args:
model_dir: Directory to download the repo to (repo files will be placed directly here)
repo_id: Hugging Face repository ID
Returns:
Path to the downloaded model if successful, None otherwise
"""
try:
print(f"Downloading BAGEL model using git lfs to {model_dir}...")
# Create parent directory if it doesn't exist
os.makedirs(model_dir, exist_ok=True)
# Check if git lfs is installed
try:
subprocess.run(["git", "lfs", "version"], check=True, capture_output=True)
except (subprocess.CalledProcessError, FileNotFoundError):
print("Git LFS not found. Installing git lfs...")
subprocess.run(["git", "lfs", "install"], check=True)
# Clone the repository directly to model_dir
clone_cmd = ["git", "clone", f"https://huggingface.co/{repo_id}", model_dir]
result = subprocess.run(clone_cmd, capture_output=True, text=True)
if result.returncode == 0:
print(f"Successfully downloaded BAGEL model to {model_dir}")
return model_dir
else:
print(f"Git clone failed: {result.stderr}")
return None
except Exception as e:
print(f"Error downloading model with git: {e}")
return None
def download_model_with_hf_hub(
model_dir: str, repo_id: str = "ByteDance-Seed/BAGEL-7B-MoT"
) -> str:
"""
Download model using huggingface_hub (fallback method)
Args:
model_dir: Directory to download the repo to (repo files will be placed directly here)
repo_id: Hugging Face repository ID
Returns:
Path to the downloaded model if successful, None otherwise
"""
try:
from huggingface_hub import snapshot_download
print(f"Downloading BAGEL model using huggingface_hub to {model_dir}...")
# Create parent directory if it doesn't exist
os.makedirs(model_dir, exist_ok=True)
# Download the entire repository directly to model_dir
snapshot_download(
repo_id=repo_id, local_dir=model_dir, local_dir_use_symlinks=False
)
print(f"Successfully downloaded BAGEL model to {model_dir}")
return model_dir
except ImportError:
print(
"huggingface_hub not installed. Please install it with: pip install huggingface_hub"
)
return None
except Exception as e:
print(f"Error downloading model with huggingface_hub: {e}")
return None
def check_model_files(model_path: str, is_df11_model: bool) -> bool:
"""
Check if core model configuration files exist.
VAE and main weights (ema.safetensors for standard) are checked separately during load.
Args:
model_path: Path to the model directory
is_df11_model: Boolean indicating if the model is DFloat11
Returns:
True if core config files exist, False otherwise
"""
required_files = [
"llm_config.json",
"vit_config.json",
]
# DFloat11 models do not have ema.safetensors in their root.
# Standard models expect ema.safetensors.
# VAE presence is checked more robustly during the loading process itself.
if not is_df11_model:
required_files.append("ema.safetensors")
for file_name in required_files:
if not os.path.exists(os.path.join(model_path, file_name)):
print(f"Missing required model file: {os.path.join(model_path, file_name)}")
return False
return True
def pil_to_tensor(img: Image.Image) -> torch.Tensor:
"""Convert PIL image to ComfyUI tensor format"""
img_array = np.array(img).astype(np.float32) / 255.0
if len(img_array.shape) == 3:
img_tensor = torch.from_numpy(img_array)[None,] # Add batch dimension
else:
img_tensor = torch.from_numpy(img_array)
return img_tensor
def tensor_to_pil(tensor: torch.Tensor) -> Image.Image:
"""Convert ComfyUI tensor to PIL image"""
if len(tensor.shape) == 4:
tensor = tensor[0] # Remove batch dimension
img_array = (tensor.cpu().numpy() * 255).astype(np.uint8)
return Image.fromarray(img_array)
def calculate_optimal_memory_gb(quantization_mode: str) -> float:
"""
Calculate optimal memory usage based on available GPU memory and quantization mode.
Args:
quantization_mode: Quantization mode (1=BF16, 2=NF4, 3=INT8)
Returns:
Optimal memory usage in GB per GPU
"""
try:
if not torch.cuda.is_available():
return 8.0 # CPU fallback
# Get total GPU memory
total_memory_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
# Reserve memory for system and other processes
if quantization_mode == "BF16": # BF16
# Full precision needs more headroom
reserved_ratio = 0.15 # Reserve 15%
optimal_ratio = 0.8 # Use 80% of remaining
elif quantization_mode == "NF4": # NF4
# 4-bit quantization is very memory efficient
reserved_ratio = 0.1 # Reserve 10%
optimal_ratio = 0.85 # Use 85% of remaining
elif quantization_mode == "INT8": # INT8
# 8-bit quantization is moderately efficient
reserved_ratio = 0.1 # Reserve 10%
optimal_ratio = 0.82 # Use 82% of remaining
else:
# Default conservative settings
reserved_ratio = 0.15
optimal_ratio = 0.75
available_memory = total_memory_gb * (1 - reserved_ratio)
optimal_memory = available_memory * optimal_ratio
# Set reasonable bounds
min_memory = 8.0
max_memory = min(80.0, total_memory_gb * 0.9) # Never exceed 90% of total
optimal_memory = max(min_memory, min(optimal_memory, max_memory))
mode_name = quantization_mode
print(
f"GPU Memory: {total_memory_gb:.1f}GB total, using {optimal_memory:.1f}GB for {mode_name} quantization"
)
return optimal_memory
except Exception as e:
print(f"Error calculating optimal memory, using default: {e}")
return 24.0
class BagelModelLoader:
"""
Unified BAGEL Model Loader with Dynamic Quantization Support
Supports both standard BAGEL and DFloat11 models with optional quantization
for ByteDance-Seed/BAGEL-7B-MoT model.
"""
SUPPORTED_MODEL_REPOS = [
"ByteDance-Seed/BAGEL-7B-MoT",
"DFloat11/BAGEL-7B-MoT-DF11",
]
QUANTIZATION_MODES = {
"BF16": "Standard (FP16/BF16)",
"NF4": "NF4 (4-bit Quantization)",
"INT8": "INT8 (8-bit Quantization)",
}
@classmethod
def INPUT_TYPES(cls):
inputs = {
"required": {
"model_repo_id": (
cls.SUPPORTED_MODEL_REPOS,
{
"default": cls.SUPPORTED_MODEL_REPOS[0],
"tooltip": "Choose BAGEL model: Standard supports quantization, DFloat11 is pre-quantized",
},
),
}
}
inputs["required"]["quantization_mode"] = (
list(cls.QUANTIZATION_MODES.keys()),
{
"default": "BF16",
"tooltip": "Quantization: BF16=Standard, NF4=4-bit, INT8=8-bit (Only for ByteDance model)",
},
)
return inputs
RETURN_TYPES = ("BAGEL_MODEL",)
RETURN_NAMES = ("model",)
FUNCTION = "load_model"
CATEGORY = "BAGEL/Core"
@classmethod
def VALIDATE_INPUTS(cls, model_repo_id, quantization_mode="BF16", **kwargs):
"""Validate input parameters"""
if model_repo_id not in cls.SUPPORTED_MODEL_REPOS:
return f"Unsupported model_repo_id: {model_repo_id}. Supported: {cls.SUPPORTED_MODEL_REPOS}"
if "DFloat11" in model_repo_id and DFloat11Model is None:
return "DFloat11 model selected, but DFloat11Model library is not installed or failed to import. Please install it: pip install dfloat11"
if quantization_mode not in cls.QUANTIZATION_MODES:
return f"Invalid quantization_mode: {quantization_mode}. Supported: {list(cls.QUANTIZATION_MODES.keys())}"
if "DFloat11" in model_repo_id and quantization_mode != "BF16":
print(
f"Warning: DFloat11 model ignores quantization_mode {quantization_mode}. Using built-in quantization."
)
if model_repo_id == "ByteDance-Seed/BAGEL-7B-MoT" and quantization_mode in [
"NF4",
"INT8",
]:
try:
import bitsandbytes as bnb
except ImportError:
return f"Quantization mode {quantization_mode} requires bitsandbytes. Please install: pip install bitsandbytes"
return True
def load_model(
self,
model_repo_id: str,
quantization_mode: str = "BF16",
) -> Tuple[Dict[str, Any]]:
"""
Load BAGEL model with unified interface supporting both standard and DFloat11 models.
Quantization is applied only to ByteDance-Seed/BAGEL-7B-MoT model.
"""
try:
is_df11_model = model_repo_id == "DFloat11/BAGEL-7B-MoT-DF11"
is_standard_model = model_repo_id == "ByteDance-Seed/BAGEL-7B-MoT"
if is_df11_model and quantization_mode != "BF16":
print(
f"DFloat11 model detected. Ignoring quantization_mode {quantization_mode} and using built-in quantization."
)
quantization_mode = "BF16"
print(
f"Loading {model_repo_id} with {self.QUANTIZATION_MODES[quantization_mode]} mode..."
)
base_repo_dir = os.path.join(comfy_models_dir, "bagel")
repo_name_segment = model_repo_id.split("/")[-1]
local_model_dir = os.path.join(base_repo_dir, repo_name_segment)
common_vae_dir = os.path.join(comfy_models_dir, "vae")
common_vae_file = os.path.join(common_vae_dir, "ae.safetensors")
if not os.path.exists(local_model_dir) or not check_model_files(
local_model_dir, is_df11_model
):
print(
f"Core model files not found or incomplete for {model_repo_id} at {local_model_dir}. Attempting download..."
)
# Attempt to download using huggingface_hub
downloaded_path = download_model_with_hf_hub(
local_model_dir, repo_id=model_repo_id
)
if not downloaded_path:
raise FileNotFoundError(
f"Failed to download BAGEL model. Please manually download it from "
f"{model_repo_id} and place it in {local_model_dir}"
)
print(f"Successfully downloaded BAGEL model to {local_model_dir}")
if not check_model_files(local_model_dir, is_df11_model):
raise FileNotFoundError(
f"Required model files missing in {local_model_dir}"
)
# Load configuration files
llm_config = Qwen2Config.from_json_file(
os.path.join(local_model_dir, "llm_config.json")
)
llm_config.qk_norm = True
llm_config.tie_word_embeddings = False
llm_config.layer_module = "Qwen2MoTDecoderLayer"
vit_config = SiglipVisionConfig.from_json_file(
os.path.join(local_model_dir, "vit_config.json")
)
vit_config.rope = False
vit_config.num_hidden_layers -= 1
vae_model, vae_config = None, None
potential_vae_paths = [
os.path.join(local_model_dir, "vae", "ae.safetensors"),
os.path.join(local_model_dir, "ae.safetensors"),
common_vae_file,
]
vae_loaded_path = None
for vae_path_to_try in potential_vae_paths:
if os.path.exists(vae_path_to_try):
try:
vae_model, vae_config = load_ae(local_path=vae_path_to_try)
if vae_model is not None and vae_config is not None:
vae_loaded_path = vae_path_to_try
break
except Exception as e:
print(f"Failed to load VAE from {vae_path_to_try}: {e}")
if not vae_loaded_path:
raise FileNotFoundError(
f"VAE model (ae.safetensors) could not be loaded from any of the expected paths: {potential_vae_paths}"
)
if is_df11_model:
config = BagelConfig(
visual_gen=True,
visual_und=True,
llm_config=llm_config,
vit_config=vit_config,
vae_config=vae_config,
vit_max_num_patch_per_side=70,
connector_act="gelu_pytorch_tanh",
latent_patch_size=2,
max_latent_size=64,
)
with init_empty_weights():
language_model = Qwen2ForCausalLM(llm_config)
vit_model = SiglipVisionModel(vit_config)
model = Bagel(language_model, vit_model, config)
model.vit_model.vision_model.embeddings.convert_conv2d_to_linear(
vit_config, meta=True
)
tokenizer = Qwen2Tokenizer.from_pretrained(local_model_dir)
tokenizer, new_token_ids, _ = add_special_tokens(tokenizer)
vae_transform = ImageTransform(1024, 512, 16)
vit_transform = ImageTransform(980, 224, 14)
model = model.to(torch.bfloat16)
model.load_state_dict(
{
name: (
torch.empty(param.shape, dtype=param.dtype, device="cpu")
if param.device.type == "meta"
else param
)
for name, param in model.state_dict().items()
},
assign=True,
)
model = DFloat11Model.from_pretrained(
local_model_dir,
bfloat16_model=model,
device="cpu",
)
# Create offload directory for model dispatch
offload_dir = os.path.join(local_model_dir, "offload")
os.makedirs(offload_dir, exist_ok=True)
device_map = infer_auto_device_map(
model,
max_memory={0: f"{calculate_optimal_memory_gb('BF16')}GiB"},
no_split_module_classes=[
"Bagel",
"Qwen2MoTDecoderLayer",
"SiglipVisionModel",
],
)
same_device_modules = [
"language_model.model.embed_tokens",
"time_embedder",
"latent_pos_embed",
"vae2llm",
"llm2vae",
"connector",
"vit_pos_embed",
]
if torch.cuda.device_count() == 1:
first_device = device_map.get(same_device_modules[0], "cuda:0")
for k in same_device_modules:
if k in device_map:
device_map[k] = first_device
else:
device_map[k] = "cuda:0"
else:
# Get first device with fallback to cuda:0
first_device = device_map.get(same_device_modules[0])
if first_device is None:
for device in device_map.values():
if isinstance(device, str) and device.startswith("cuda"):
first_device = device
break
else:
first_device = "cuda:0"
for k in same_device_modules:
device_map[k] = first_device
model = dispatch_model(
model,
device_map=device_map,
offload_dir=offload_dir,
force_hooks=True,
)
model = model.eval()
inferencer = InterleaveInferencer(
model=model,
vae_model=vae_model,
tokenizer=tokenizer,
vae_transform=vae_transform,
vit_transform=vit_transform,
new_token_ids=new_token_ids,
)
model_dict = {
"model": model,
"inferencer": inferencer,
"tokenizer": tokenizer,
"vae_model": vae_model,
"vae_transform": vae_transform,
"vit_transform": vit_transform,
"config": config,
"model_path": local_model_dir,
"model_repo_id": model_repo_id,
"is_df11": is_df11_model,
"device": "cuda" if torch.cuda.is_available() else "cpu",
}
print(f"Successfully loaded BAGEL DF11 model from {local_model_dir}")
return (model_dict,)
elif is_standard_model:
# Create BAGEL configuration
config = BagelConfig(
visual_gen=True,
visual_und=True,
llm_config=llm_config,
vit_config=vit_config,
vae_config=vae_config,
vit_max_num_patch_per_side=70,
connector_act="gelu_pytorch_tanh",
latent_patch_size=2,
max_latent_size=64,
)
# Initialize empty model
with init_empty_weights():
language_model = Qwen2ForCausalLM(llm_config)
vit_model = SiglipVisionModel(vit_config)
model = Bagel(language_model, vit_model, config)
model.vit_model.vision_model.embeddings.convert_conv2d_to_linear(
vit_config, meta=True
)
tokenizer = Qwen2Tokenizer.from_pretrained(local_model_dir)
tokenizer, new_token_ids, _ = add_special_tokens(tokenizer)
vae_transform = ImageTransform(1024, 512, 16)
vit_transform = ImageTransform(980, 224, 14)
if quantization_mode in ["NF4", "INT8"]:
max_memory_gb = calculate_optimal_memory_gb(quantization_mode)
max_memory_per_gpu = f"{max_memory_gb}GiB"
device_map = infer_auto_device_map(
model,
max_memory={
i: max_memory_per_gpu
for i in range(torch.cuda.device_count())
},
no_split_module_classes=["Bagel", "Qwen2MoTDecoderLayer"],
)
else:
device_map = infer_auto_device_map(
model,
no_split_module_classes=["Bagel", "Qwen2MoTDecoderLayer"],
)
# Ensure same device placement for critical modules
same_device_modules = [
"language_model.model.embed_tokens",
"time_embedder",
"latent_pos_embed",
"vae2llm",
"llm2vae",
"connector",
"vit_pos_embed",
]
if torch.cuda.device_count() == 1:
first_device = device_map.get(same_device_modules[0], "cuda:0")
for k in same_device_modules:
if k in device_map:
device_map[k] = first_device
else:
device_map[k] = "cuda:0"
else:
first_device = device_map.get(same_device_modules[0])
if first_device is None:
for device in device_map.values():
if isinstance(device, str) and device.startswith("cuda"):
first_device = device
break
else:
first_device = "cuda:0"
for k in same_device_modules:
if k in device_map:
device_map[k] = first_device
else:
device_map[k] = first_device
# Load model based on quantization mode
checkpoint_path = os.path.join(local_model_dir, "ema.safetensors")
if quantization_mode == "BF16":
# Standard loading
print("Loading model in standard FP16/BF16 mode...")
model = load_checkpoint_and_dispatch(
model,
checkpoint=checkpoint_path,
device_map=device_map,
offload_buffers=True,
offload_folder="offload",
dtype=torch.bfloat16,
force_hooks=True,
).eval()
elif quantization_mode == "NF4":
print("Loading model with NF4 (4-bit) quantization...")
bnb_quantization_config = BnbQuantizationConfig(
load_in_4bit=True,
bnb_4bit_compute_dtype=torch.bfloat16,
bnb_4bit_use_double_quant=False,
bnb_4bit_quant_type="nf4",
)
model = load_and_quantize_model(
model,
weights_location=checkpoint_path,
bnb_quantization_config=bnb_quantization_config,
device_map=device_map,
offload_folder="offload",
).eval()
elif quantization_mode == "INT8":
print("Loading model with INT8 (8-bit) quantization...")
bnb_quantization_config = BnbQuantizationConfig(
load_in_8bit=True, torch_dtype=torch.bfloat16
)
model = load_and_quantize_model(
model,
weights_location=checkpoint_path,
bnb_quantization_config=bnb_quantization_config,
device_map=device_map,
offload_folder="offload",
).eval()
else:
raise ValueError(
f"Unsupported quantization mode: {quantization_mode}"
)
# Create inferencer
inferencer = InterleaveInferencer(
model=model,
vae_model=vae_model,
tokenizer=tokenizer,
vae_transform=vae_transform,
vit_transform=vit_transform,
new_token_ids=new_token_ids,
)
model_dict = {
"model": model,
"inferencer": inferencer,
"tokenizer": tokenizer,
"vae_model": vae_model,
"vae_transform": vae_transform,
"vit_transform": vit_transform,
"config": config,
"model_path": local_model_dir,
"model_repo_id": model_repo_id,
"quantization_mode": quantization_mode,
"quantization_info": self.QUANTIZATION_MODES[quantization_mode],
"is_df11": False,
"device": "cuda" if torch.cuda.is_available() else "cpu",
}
print(
f"Successfully loaded BAGEL model with {self.QUANTIZATION_MODES[quantization_mode]} mode"
)
return (model_dict,)
except Exception as e:
print(f"Error loading BAGEL model: {e}")
raise e
class BagelTextToImage:
"""BAGEL Text to Image Node"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("BAGEL_MODEL", {"tooltip": "BAGEL model"}),
"prompt": (
"STRING",
{
"multiline": True,
"default": "A female cosplayer portraying an ethereal fairy or elf, wearing a flowing dress made of delicate fabrics in soft, mystical colors like emerald green and silver. She has pointed ears, a gentle, enchanting expression, and her outfit is adorned with sparkling jewels and intricate patterns. The background is a magical forest with glowing plants, mystical creatures, and a serene atmosphere.",
"tooltip": "Text prompt",
},
),
"seed": (
"INT",
{
"default": 0,
"min": 0,
"max": 1000000,
"tooltip": "Random seed, 0 for random",
},
),
"image_ratio": (
["1:1", "4:3", "3:4", "16:9", "9:16"],
{"default": "1:1", "tooltip": "Image aspect ratio"},
),
"cfg_text_scale": (
"FLOAT",
{
"default": 4.0,
"min": 1.0,
"max": 8.0,
"step": 0.1,
"tooltip": "CFG text scaling",
},
),
"num_timesteps": (
"INT",
{
"default": 50,
"min": 10,
"max": 100,
"step": 5,
"tooltip": "Denoising steps",
},
),
},
"optional": {
"show_thinking": (
"BOOLEAN",
{"default": False, "tooltip": "Display reasoning process"},
),
"cfg_interval": (
"FLOAT",
{
"default": 0.4,
"min": 0.0,
"max": 1.0,
"step": 0.1,
"tooltip": "CFG interval start value",
},
),
"timestep_shift": (
"FLOAT",
{
"default": 3.0,
"min": 1.0,
"max": 5.0,
"step": 0.5,
"tooltip": "Timestep offset",
},
),
"cfg_renorm_min": (
"FLOAT",
{
"default": 0.0,
"min": 0.0,
"max": 1.0,
"step": 0.1,
"tooltip": "CFG re-normalization minimum value",
},
),
"cfg_renorm_type": (
["global", "local", "text_channel"],
{"default": "global", "tooltip": "CFG re-normalization type"},
),
"text_temperature": (
"FLOAT",
{
"default": 0.3,
"min": 0.0,
"max": 1.0,
"step": 0.1,
"tooltip": "Text generation temperature",
},
),
},
}
RETURN_TYPES = ("IMAGE", "STRING")
RETURN_NAMES = ("image", "thinking")
FUNCTION = "generate_image"
CATEGORY = "BAGEL/Core"
@classmethod
def VALIDATE_INPUTS(
cls, model, prompt, seed, image_ratio, cfg_text_scale, num_timesteps, **kwargs
):
"""Validate input parameters"""
if not isinstance(prompt, str) or not prompt.strip():
return "Prompt must be a non-empty string"
if not isinstance(seed, int) or seed < 0:
return "Seed must be a non-negative integer"
if image_ratio not in ["1:1", "4:3", "3:4", "16:9", "9:16"]:
return f"Invalid image_ratio: {image_ratio}"
if (
not isinstance(cfg_text_scale, (int, float))
or cfg_text_scale < 1.0
or cfg_text_scale > 8.0
):
return "cfg_text_scale must be between 1.0 and 8.0"
if (
not isinstance(num_timesteps, int)
or num_timesteps < 10
or num_timesteps > 100
):
return "num_timesteps must be between 10 and 100"
return True
def generate_image(
self,
model: Dict[str, Any],
prompt: str,
seed: int,
image_ratio: str,
cfg_text_scale: float,
num_timesteps: int,
show_thinking: bool = False,
cfg_interval: float = 0.4,
timestep_shift: float = 3.0,
cfg_renorm_min: float = 0.0,
cfg_renorm_type: str = "global",
text_temperature: float = 0.3,
) -> Tuple[torch.Tensor, str]:
"""
Generate image from text using BAGEL model
Args:
model: BAGEL model dictionary
prompt: Text prompt
seed: Random seed
image_ratio: Image aspect ratio
cfg_text_scale: CFG text scaling
num_timesteps: Denoising steps
show_thinking: Whether to display the reasoning process
cfg_interval: CFG interval start value
timestep_shift: Timestep offset
cfg_renorm_min: CFG re-normalization minimum value
cfg_renorm_type: CFG re-normalization type
text_temperature: Text generation temperature
Returns:
Generated image tensor and reasoning process text
"""
try:
# Set random seed
set_seed(seed)
# Get inferencer
inferencer = model["inferencer"]
# Set image dimensions
image_shapes_map = {
"1:1": (1024, 1024),
"4:3": (768, 1024),
"3:4": (1024, 768),
"16:9": (576, 1024),
"9:16": (1024, 576),
}
image_shapes = image_shapes_map[image_ratio]
# Set inference hyperparameters
inference_hyper = {
"max_think_token_n": 1024 if show_thinking else 1024,
"do_sample": False if not show_thinking else False,
"text_temperature": text_temperature if show_thinking else 0.3,
"cfg_text_scale": cfg_text_scale,
"cfg_interval": [cfg_interval, 1.0], # End value fixed at 1.0
"timestep_shift": timestep_shift,
"num_timesteps": num_timesteps,
"cfg_renorm_min": cfg_renorm_min,
"cfg_renorm_type": cfg_renorm_type,
"image_shapes": image_shapes,
}
# Initialize ProgressBar
# The loop in Bagel.generate_image runs for (num_timesteps - 1) iterations
actual_iterations = num_timesteps - 1 if num_timesteps > 0 else 0
pbar = ProgressBar(actual_iterations)
# Call inferencer, passing pbar
result = inferencer(
text=prompt, think=show_thinking, pbar=pbar, **inference_hyper
)
# Convert image format
pil_image = result["image"]
tensor_image = pil_to_tensor(pil_image)
# Get reasoning process
thinking_text = result.get("text", "") if show_thinking else ""
print(f"Generated image with size: {pil_image.size}")
return (tensor_image, thinking_text)
except Exception as e:
print(f"Error in text to image generation: {e}")
# Return empty image and error message
empty_image = torch.zeros((1, 512, 512, 3))
return (empty_image, f"Error: {str(e)}")
class BagelImageEdit:
"""BAGEL Image Edit Node"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("BAGEL_MODEL", {"tooltip": "BAGEL model"}),
"image": ("IMAGE", {"tooltip": "Input image"}),
"prompt": (
"STRING",
{
"multiline": True,
"default": "Edit the image according to the description",
"tooltip": "Editing prompt",
},
),
"seed": (
"INT",
{
"default": 0,
"min": 0,
"max": 1000000,
"tooltip": "Random seed, 0 for random",
},
),
"cfg_text_scale": (
"FLOAT",
{
"default": 4.0,
"min": 1.0,
"max": 8.0,
"step": 0.1,
"tooltip": "CFG text scaling",
},
),
"cfg_img_scale": (
"FLOAT",
{
"default": 2.0,
"min": 1.0,
"max": 4.0,
"step": 0.1,
"tooltip": "CFG image scaling",
},
),
"num_timesteps": (
"INT",
{
"default": 50,
"min": 10,
"max": 100,
"step": 5,
"tooltip": "Denoising steps",
},
),
},
"optional": {
"show_thinking": (
"BOOLEAN",
{"default": False, "tooltip": "Display reasoning process"},
),
"cfg_interval": (
"FLOAT",
{
"default": 0.0,
"min": 0.0,
"max": 1.0,
"step": 0.1,
"tooltip": "CFG interval start value",
},
),
"timestep_shift": (
"FLOAT",
{
"default": 3.0,
"min": 1.0,
"max": 10.0,
"step": 0.5,
"tooltip": "Timestep offset",
},
),
"cfg_renorm_min": (
"FLOAT",
{
"default": 0.0,
"min": 0.0,
"max": 1.0,
"step": 0.1,
"tooltip": "CFG re-normalization minimum value",
},
),
"cfg_renorm_type": (
["global", "local", "text_channel"],
{"default": "text_channel", "tooltip": "CFG re-normalization type"},
),
"text_temperature": (
"FLOAT",
{
"default": 0.3,
"min": 0.0,
"max": 1.0,
"step": 0.1,
"tooltip": "Text generation temperature",
},
),
},
}
RETURN_TYPES = ("IMAGE", "STRING")
RETURN_NAMES = ("image", "thinking")
FUNCTION = "edit_image"
CATEGORY = "BAGEL/Core"
@classmethod
def VALIDATE_INPUTS(
cls,
model,
image,
prompt,
seed,
cfg_text_scale,
cfg_img_scale,
num_timesteps,
**kwargs,
):
"""Validate input parameters"""
if not isinstance(prompt, str) or not prompt.strip():
return "Prompt must be a non-empty string"
if not isinstance(seed, int) or seed < 0:
return "Seed must be a non-negative integer"
if (
not isinstance(cfg_text_scale, (int, float))
or cfg_text_scale < 1.0
or cfg_text_scale > 8.0
):
return "cfg_text_scale must be between 1.0 and 8.0"
if (
not isinstance(cfg_img_scale, (int, float))
or cfg_img_scale < 1.0
or cfg_img_scale > 4.0
):
return "cfg_img_scale must be between 1.0 and 4.0"
if (
not isinstance(num_timesteps, int)
or num_timesteps < 10
or num_timesteps > 100
):
return "num_timesteps must be between 10 and 100"
return True
def edit_image(
self,
model: Dict[str, Any],
image: torch.Tensor,
prompt: str,
seed: int,
cfg_text_scale: float,
cfg_img_scale: float,
num_timesteps: int,
show_thinking: bool = False,
cfg_interval: float = 0.0,
timestep_shift: float = 3.0,
cfg_renorm_min: float = 0.0,
cfg_renorm_type: str = "text_channel",
text_temperature: float = 0.3,
) -> Tuple[torch.Tensor, str]:
"""
Edit image using BAGEL model
Args:
model: BAGEL model dictionary
image: Input image tensor
prompt: Editing prompt
seed: Random seed
cfg_text_scale: CFG text scaling
cfg_img_scale: CFG image scaling
num_timesteps: Denoising steps
show_thinking: Whether to display the reasoning process
cfg_interval: CFG interval start value
timestep_shift: Timestep offset
cfg_renorm_min: CFG re-normalization minimum value
cfg_renorm_type: CFG re-normalization type
text_temperature: Text generation temperature
Returns:
Edited image tensor and reasoning process text
"""
try:
# Set random seed
set_seed(seed)
# Get inferencer
inferencer = model["inferencer"]
# Convert image format
pil_image = tensor_to_pil(image)
pil_image = pil_img2rgb(pil_image)
# Set inference hyperparameters
inference_hyper = {
"max_think_token_n": 1024 if show_thinking else 1024,
"do_sample": False if not show_thinking else False,
"text_temperature": text_temperature if show_thinking else 0.3,
"cfg_text_scale": cfg_text_scale,
"cfg_img_scale": cfg_img_scale,
"cfg_interval": [cfg_interval, 1.0], # End value fixed at 1.0
"timestep_shift": timestep_shift,
"num_timesteps": num_timesteps,
"cfg_renorm_min": cfg_renorm_min,
"cfg_renorm_type": cfg_renorm_type,
}
# Initialize ProgressBar
actual_iterations = num_timesteps - 1 if num_timesteps > 0 else 0
pbar = ProgressBar(actual_iterations)
# Call inferencer, passing pbar
result = inferencer(
image=pil_image,
text=prompt,
think=show_thinking,
pbar=pbar,
**inference_hyper,
)
# Convert image format
edited_pil_image = result["image"]
tensor_image = pil_to_tensor(edited_pil_image)
# Get reasoning process
thinking_text = result.get("text", "") if show_thinking else ""
print(f"Edited image with size: {edited_pil_image.size}")
return (tensor_image, thinking_text)
except Exception as e:
print(f"Error in image editing: {e}")
# Return original image and error message
return (image, f"Error: {str(e)}")
class BagelImageUnderstanding:
"""BAGEL Image Understanding Node"""
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"model": ("BAGEL_MODEL", {"tooltip": "BAGEL model"}),
"image": ("IMAGE", {"tooltip": "Input image"}),
"prompt": (
"STRING",
{
"multiline": True,
"default": "What do you see in this image?",
"tooltip": "Question text",
},
),
},
"optional": {
"show_thinking": (
"BOOLEAN",
{"default": False, "tooltip": "Display reasoning process"},
),
"do_sample": (
"BOOLEAN",
{"default": False, "tooltip": "Enable sampling"},
),
"text_temperature": (
"FLOAT",
{
"default": 0.3,
"min": 0.0,
"max": 1.0,
"step": 0.05,
"tooltip": "Text generation temperature",
},
),
"max_new_tokens": (
"INT",
{
"default": 512,
"min": 64,
"max": 4096,
"step": 64,
"tooltip": "Maximum new tokens",
},
),
},
}
RETURN_TYPES = ("STRING",)
RETURN_NAMES = ("text",)
FUNCTION = "understand_image"
CATEGORY = "BAGEL/Core"
@classmethod
def VALIDATE_INPUTS(cls, model, image, prompt, **kwargs):
"""Validate input parameters"""
if not isinstance(prompt, str) or not prompt.strip():
return "Prompt must be a non-empty string"
# Validate optional parameters
if "text_temperature" in kwargs:
temp = kwargs["text_temperature"]
if not isinstance(temp, (int, float)) or temp < 0.0 or temp > 1.0:
return "text_temperature must be between 0.0 and 1.0"
if "max_new_tokens" in kwargs:
tokens = kwargs["max_new_tokens"]
if not isinstance(tokens, int) or tokens < 64 or tokens > 4096:
return "max_new_tokens must be between 64 and 4096"
return True
def understand_image(
self,
model: Dict[str, Any],
image: torch.Tensor,
prompt: str,
show_thinking: bool = False,
do_sample: bool = False,
text_temperature: float = 0.3,
max_new_tokens: int = 512,
) -> Tuple[str]:
"""
Use BAGEL model to understand image and answer questions
Args:
model: BAGEL model dictionary
image: Input image tensor
prompt: Question text
show_thinking: Whether to display the reasoning process
do_sample: Whether to enable sampling
text_temperature: Text generation temperature
max_new_tokens: Maximum new tokens
Returns:
Answer text
"""
try:
# Get inferencer
inferencer = model["inferencer"]
# Convert image format
pil_image = tensor_to_pil(image)
pil_image = pil_img2rgb(pil_image)
# Set inference hyperparameters
inference_hyper = {
"do_sample": do_sample,
"text_temperature": text_temperature,
"max_think_token_n": max_new_tokens,
}
# Call inferencer
result = inferencer(
image=pil_image,
text=prompt,
think=show_thinking,
understanding_output=True,
**inference_hyper,
)
answer_text = result["text"]
print(f"Image understanding completed, response length: {len(answer_text)}")
return (answer_text,)
except Exception as e:
print(f"Error in image understanding: {e}")
return (f"Error: {str(e)}",)
# Node mappings
NODE_CLASS_MAPPINGS = {
"BagelModelLoader": BagelModelLoader,
"BagelTextToImage": BagelTextToImage,
"BagelImageEdit": BagelImageEdit,
"BagelImageUnderstanding": BagelImageUnderstanding,
}
# Display name mappings
NODE_DISPLAY_NAME_MAPPINGS = {
"BagelModelLoader": "BAGEL Model Loader",
"BagelTextToImage": "BAGEL Text to Image",
"BagelImageEdit": "BAGEL Image Edit",
"BagelImageUnderstanding": "BAGEL Image Understanding",
}
# Export for ComfyUI
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]