feat: add comfyui sdk
This commit is contained in:
193
cargos/comfyui-sdk/templates/workflow_template.rs
Normal file
193
cargos/comfyui-sdk/templates/workflow_template.rs
Normal file
@@ -0,0 +1,193 @@
|
||||
//! WorkflowTemplate implementation for managing parameterized ComfyUI workflows
|
||||
|
||||
use std::collections::HashMap;
|
||||
use crate::types::{
|
||||
WorkflowTemplateData, TemplateMetadata, ComfyUIWorkflow,
|
||||
ParameterSchema, ParameterValues, ValidationResult, ParameterType
|
||||
};
|
||||
use crate::error::{ComfyUIError, Result};
|
||||
use crate::templates::WorkflowInstance;
|
||||
use crate::utils::validation::validate_parameters;
|
||||
|
||||
/// Represents a parameterized workflow template
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct WorkflowTemplate {
|
||||
metadata: TemplateMetadata,
|
||||
workflow: ComfyUIWorkflow,
|
||||
parameters: HashMap<String, ParameterSchema>,
|
||||
}
|
||||
|
||||
impl WorkflowTemplate {
|
||||
/// Creates a new WorkflowTemplate instance
|
||||
pub fn new(data: WorkflowTemplateData) -> Result<Self> {
|
||||
let template = Self {
|
||||
metadata: data.metadata,
|
||||
workflow: data.workflow,
|
||||
parameters: data.parameters,
|
||||
};
|
||||
|
||||
// Validate template structure
|
||||
template.validate_template()?;
|
||||
Ok(template)
|
||||
}
|
||||
|
||||
/// Gets the template metadata
|
||||
pub fn metadata(&self) -> &TemplateMetadata {
|
||||
&self.metadata
|
||||
}
|
||||
|
||||
/// Gets the workflow definition
|
||||
pub fn workflow(&self) -> &ComfyUIWorkflow {
|
||||
&self.workflow
|
||||
}
|
||||
|
||||
/// Gets the parameter schemas
|
||||
pub fn parameters(&self) -> &HashMap<String, ParameterSchema> {
|
||||
&self.parameters
|
||||
}
|
||||
|
||||
/// Gets the template ID
|
||||
pub fn id(&self) -> &str {
|
||||
&self.metadata.id
|
||||
}
|
||||
|
||||
/// Gets the template name
|
||||
pub fn name(&self) -> &str {
|
||||
&self.metadata.name
|
||||
}
|
||||
|
||||
/// Gets the template description
|
||||
pub fn description(&self) -> Option<&str> {
|
||||
self.metadata.description.as_deref()
|
||||
}
|
||||
|
||||
/// Gets the template category
|
||||
pub fn category(&self) -> Option<&str> {
|
||||
self.metadata.category.as_deref()
|
||||
}
|
||||
|
||||
/// Gets the template tags
|
||||
pub fn tags(&self) -> &[String] {
|
||||
self.metadata.tags.as_deref().unwrap_or(&[])
|
||||
}
|
||||
|
||||
/// Validates the provided parameters against the template schema
|
||||
pub fn validate(&self, parameters: &ParameterValues) -> ValidationResult {
|
||||
validate_parameters(parameters, &self.parameters)
|
||||
}
|
||||
|
||||
/// Creates a workflow instance with the provided parameters
|
||||
pub fn create_instance(&self, parameters: ParameterValues) -> Result<WorkflowInstance> {
|
||||
let validation = self.validate(¶meters);
|
||||
if !validation.valid {
|
||||
let error_messages: Vec<String> = validation.errors
|
||||
.iter()
|
||||
.map(|e| format!("{}: {}", e.path, e.message))
|
||||
.collect();
|
||||
return Err(ComfyUIError::parameter_validation(
|
||||
format!("Invalid parameters: {}", error_messages.join(", "))
|
||||
));
|
||||
}
|
||||
|
||||
WorkflowInstance::new(self.clone(), parameters)
|
||||
}
|
||||
|
||||
/// Gets the parameter schema for a specific parameter
|
||||
pub fn get_parameter_schema(&self, parameter_name: &str) -> Option<&ParameterSchema> {
|
||||
self.parameters.get(parameter_name)
|
||||
}
|
||||
|
||||
/// Gets all required parameter names
|
||||
pub fn get_required_parameters(&self) -> Vec<&str> {
|
||||
self.parameters
|
||||
.iter()
|
||||
.filter(|(_, schema)| schema.required.unwrap_or(false))
|
||||
.map(|(name, _)| name.as_str())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Gets all optional parameter names
|
||||
pub fn get_optional_parameters(&self) -> Vec<&str> {
|
||||
self.parameters
|
||||
.iter()
|
||||
.filter(|(_, schema)| !schema.required.unwrap_or(false))
|
||||
.map(|(name, _)| name.as_str())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Checks if the template has a specific parameter
|
||||
pub fn has_parameter(&self, parameter_name: &str) -> bool {
|
||||
self.parameters.contains_key(parameter_name)
|
||||
}
|
||||
|
||||
/// Gets default values for all parameters that have defaults
|
||||
pub fn get_default_values(&self) -> ParameterValues {
|
||||
let mut defaults = HashMap::new();
|
||||
for (name, schema) in &self.parameters {
|
||||
if let Some(default_value) = &schema.default {
|
||||
defaults.insert(name.clone(), default_value.clone());
|
||||
}
|
||||
}
|
||||
defaults
|
||||
}
|
||||
|
||||
/// Creates a copy of the template with updated metadata
|
||||
pub fn with_metadata(&self, metadata: TemplateMetadata) -> Self {
|
||||
Self {
|
||||
metadata,
|
||||
workflow: self.workflow.clone(),
|
||||
parameters: self.parameters.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Converts the template to a JSON-serializable object
|
||||
pub fn to_data(&self) -> WorkflowTemplateData {
|
||||
WorkflowTemplateData {
|
||||
metadata: self.metadata.clone(),
|
||||
workflow: self.workflow.clone(),
|
||||
parameters: self.parameters.clone(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Creates a WorkflowTemplate from template data
|
||||
pub fn from_data(data: WorkflowTemplateData) -> Result<Self> {
|
||||
Self::new(data)
|
||||
}
|
||||
|
||||
/// Validates the template structure
|
||||
fn validate_template(&self) -> Result<()> {
|
||||
// Validate metadata
|
||||
if self.metadata.id.is_empty() {
|
||||
return Err(ComfyUIError::template_validation(
|
||||
"Template metadata must have a valid id"
|
||||
));
|
||||
}
|
||||
if self.metadata.name.is_empty() {
|
||||
return Err(ComfyUIError::template_validation(
|
||||
"Template metadata must have a valid name"
|
||||
));
|
||||
}
|
||||
|
||||
// Validate workflow
|
||||
if self.workflow.is_empty() {
|
||||
return Err(ComfyUIError::template_validation(
|
||||
"Template must have a valid workflow object"
|
||||
));
|
||||
}
|
||||
|
||||
// Validate parameter schemas
|
||||
for (name, schema) in &self.parameters {
|
||||
if !matches!(
|
||||
schema.param_type,
|
||||
ParameterType::String | ParameterType::Number | ParameterType::Boolean |
|
||||
ParameterType::Array | ParameterType::Object
|
||||
) {
|
||||
return Err(ComfyUIError::template_validation(
|
||||
format!("Parameter '{}' has invalid type: {:?}", name, schema.param_type)
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user