246 lines
9.7 KiB
Python
246 lines
9.7 KiB
Python
"""
|
|
WorkflowGenerator: Generates workflow diagrams from natural language descriptions.
|
|
Enhanced with GPU layer management, error handling, and ComfyUI integration.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import time
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
from ..utils.model_loaders.device_utils import load_model
|
|
from ..utils.process_utils import del_digram_primitive
|
|
|
|
|
|
class WorkflowGenerator:
|
|
"""
|
|
Generator responsible for creating workflow diagrams from natural language.
|
|
|
|
Uses a language model to convert user descriptions into structured workflow
|
|
diagrams represented as JSON edge lists.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
cfg: Any,
|
|
n_gpu_layers: int | str = "auto",
|
|
dtype: str | torch.dtype = "auto",
|
|
device_preference: str = "auto",
|
|
attn_implementation: str = "auto",
|
|
context_size: int = 4096,
|
|
use_mmap: bool = True,
|
|
use_mlock: bool = False,
|
|
n_batch: int = 512,
|
|
n_threads: int | None = None,
|
|
):
|
|
"""
|
|
Initialize WorkflowGenerator with model configuration.
|
|
|
|
Args:
|
|
cfg: Configuration object with workflow_generator settings
|
|
n_gpu_layers: For GGUF: "auto", -1 for all, 0 for CPU, or specific count
|
|
dtype: For HuggingFace: "auto", "fp16", "bf16", "fp32", "fp8", or torch.dtype
|
|
device_preference: "auto", "cuda", or "cpu"
|
|
attn_implementation: For HuggingFace: "auto", "flash_attention_2", or "sdpa"
|
|
"""
|
|
logging.info("WorkflowGenerator: Initializing...")
|
|
|
|
# Get model format from config (default to huggingface for backward compatibility)
|
|
model_format = getattr(cfg.workflow_generator, "model_format", "huggingface")
|
|
model_path = cfg.workflow_generator.model_path
|
|
tokenizer_path = getattr(cfg.workflow_generator, "tokenizer_path", None)
|
|
|
|
logging.info(f"Model format: {model_format}, path: {model_path}")
|
|
|
|
load_kwargs = {}
|
|
|
|
if model_format.lower() == "huggingface":
|
|
# HuggingFace models use dtype, device_map, and attention implementation
|
|
load_kwargs = {
|
|
"dtype": dtype,
|
|
"device_preference": device_preference,
|
|
"attn_implementation": attn_implementation,
|
|
}
|
|
logging.debug("Model format: HuggingFace")
|
|
else:
|
|
# GGUF models use n_gpu_layers, context_size, and performance parameters
|
|
load_kwargs = {
|
|
"n_gpu_layers": n_gpu_layers,
|
|
"device_preference": device_preference,
|
|
"context_size": context_size,
|
|
"use_mmap": use_mmap,
|
|
"use_mlock": use_mlock,
|
|
"n_batch": n_batch,
|
|
"n_threads": n_threads,
|
|
}
|
|
logging.debug("Model format: GGUF (already quantized)")
|
|
|
|
load_start = time.time()
|
|
try:
|
|
self.model_wrapper = load_model(model_format=model_format, model_path=model_path, tokenizer_path=tokenizer_path, **load_kwargs)
|
|
load_time = time.time() - load_start
|
|
logging.info(f"Model loaded in {load_time:.2f}s")
|
|
except Exception as e:
|
|
logging.error(f"Failed to load model: {e}")
|
|
raise
|
|
|
|
self.model = self.model_wrapper.model
|
|
self.tokenizer = self.model_wrapper.tokenizer
|
|
|
|
try:
|
|
if torch.cuda.is_available() and model_format.lower() == "huggingface":
|
|
vram_used = torch.cuda.memory_allocated() / 1024**3
|
|
logging.debug(f"GPU memory allocated: {vram_used:.2f} GB")
|
|
except Exception:
|
|
pass
|
|
|
|
self.top_p = getattr(cfg.workflow_generator, "top_p", 0.7)
|
|
self.max_new_tokens = getattr(cfg.workflow_generator, "max_new_tokens", 8192)
|
|
self.temperature = getattr(cfg.workflow_generator, "temperature", 0.95)
|
|
self.allow_primitive_nodes = getattr(cfg.workflow_generator, "allow_primitive_nodes", False)
|
|
# Template matches FlowDataset training format: instruction + "description: [description]"
|
|
# Model (Qwen2.5-14B, not instruct) was fine-tuned on this exact format
|
|
self.template = "Based on the description I provided, generate a JSON example of the required ComfyUi workflow. description:"
|
|
|
|
logging.info("WorkflowGenerator: Initialized successfully")
|
|
|
|
def generate(self, prompt: str, seed: int | None = None) -> list:
|
|
"""
|
|
Generate workflow diagram from natural language prompt.
|
|
|
|
Args:
|
|
prompt: Natural language description of the desired workflow
|
|
seed: Random seed for reproducible generation (optional)
|
|
|
|
Returns:
|
|
List of edges representing the workflow diagram
|
|
"""
|
|
# Format prompt to match FlowDataset training format: Model (Qwen2.5-14B, not instruct) was fine-tuned on this exact format
|
|
connection_guidance = (
|
|
" Ensure all connections are valid: every input_name must exist on the destination node, "
|
|
"and every output_name must exist on the source node."
|
|
)
|
|
|
|
if "Based on the description I provided" not in prompt:
|
|
formatted_prompt = self.template + " " + prompt + connection_guidance
|
|
else:
|
|
if "description:" in prompt:
|
|
desc_idx = prompt.find("description:")
|
|
if desc_idx != -1:
|
|
desc_prefix = prompt[: desc_idx + len("description:")]
|
|
desc_content = prompt[desc_idx + len("description:") :].strip()
|
|
formatted_prompt = desc_prefix + " " + desc_content + connection_guidance
|
|
else:
|
|
formatted_prompt = prompt + connection_guidance
|
|
else:
|
|
formatted_prompt = prompt + " description: " + connection_guidance
|
|
|
|
prompt = formatted_prompt
|
|
|
|
messages = [{"role": "user", "content": prompt}]
|
|
|
|
try:
|
|
# Use HuggingFace tokenizer even for GGUF models (embedded tokenizer doesn't expose apply_chat_template)
|
|
text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
|
except Exception as e:
|
|
logging.error(f"Failed to apply chat template: {e}")
|
|
text = prompt
|
|
|
|
logging.debug(f"Tokenizing input (max_new_tokens={self.max_new_tokens}, temperature={self.temperature}, top_p={self.top_p})...")
|
|
|
|
try:
|
|
model_inputs = self.tokenizer([text], return_tensors="pt")
|
|
if hasattr(self.model_wrapper, "device") and isinstance(self.model_wrapper.device, torch.device):
|
|
model_inputs = model_inputs.to(self.model_wrapper.device)
|
|
input_tokens = model_inputs.input_ids.shape[1]
|
|
logging.debug(f"Input tokens: {input_tokens}")
|
|
except Exception as e:
|
|
logging.error(f"Failed to tokenize input: {e}")
|
|
raise
|
|
|
|
logging.info("Generating workflow diagram...")
|
|
|
|
try:
|
|
generate_kwargs = {
|
|
"max_new_tokens": self.max_new_tokens,
|
|
"temperature": self.temperature,
|
|
"do_sample": True,
|
|
"top_p": self.top_p,
|
|
}
|
|
if seed is not None:
|
|
max_seed = 2**32 - 1
|
|
if seed < 0 or seed > max_seed:
|
|
seed = seed % (max_seed + 1)
|
|
generate_kwargs["seed"] = seed
|
|
|
|
generated_ids = self.model_wrapper.generate(model_inputs.input_ids, **generate_kwargs)
|
|
|
|
generated_ids = [output_ids[len(input_ids) :] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)]
|
|
output_tokens = len(generated_ids[0])
|
|
logging.debug(f"Generated {output_tokens} tokens")
|
|
except Exception as e:
|
|
logging.error(f"Failed to generate tokens: {e}")
|
|
raise
|
|
|
|
try:
|
|
response = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
|
|
except Exception as e:
|
|
logging.error(f"Failed to decode tokens: {e}")
|
|
raise
|
|
|
|
logging.debug("Post-processing diagram...")
|
|
try:
|
|
diagram = self.post_process(response)
|
|
logging.info(f"Workflow diagram generated successfully: {len(diagram)} edges")
|
|
return diagram
|
|
except Exception as e:
|
|
logging.error(f"Failed to post-process diagram: {e}")
|
|
raise
|
|
|
|
def post_process(self, response: str | list) -> list:
|
|
"""
|
|
Post-process the model response into a valid diagram format.
|
|
|
|
Args:
|
|
response: Raw model response (string or list)
|
|
|
|
Returns:
|
|
List of edges representing the workflow diagram
|
|
"""
|
|
diagram = response
|
|
|
|
if isinstance(diagram, str):
|
|
try:
|
|
diagram = json.loads(diagram)
|
|
except json.JSONDecodeError as e:
|
|
logging.error(f"Failed to parse JSON response: {e}")
|
|
logging.error(f"Response content: {response[:500]}")
|
|
raise
|
|
|
|
if not self.allow_primitive_nodes:
|
|
try:
|
|
diagram = del_digram_primitive(diagram)
|
|
except Exception as e:
|
|
logging.warning(f"Failed to remove primitive nodes: {e}")
|
|
|
|
if not isinstance(diagram, list):
|
|
raise ValueError(f"Invalid diagram format: expected list, got {type(diagram)}")
|
|
|
|
return diagram
|
|
|
|
def cleanup(self):
|
|
"""Clean up model resources."""
|
|
if hasattr(self, "model_wrapper"):
|
|
del self.model_wrapper
|
|
if hasattr(self, "model"):
|
|
del self.model
|
|
if hasattr(self, "tokenizer"):
|
|
del self.tokenizer
|
|
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
logging.info("WorkflowGenerator: Resources cleaned up")
|