feat: add comfyui sdk
This commit is contained in:
194
cargos/comfyui-sdk/templates/template_manager.rs
Normal file
194
cargos/comfyui-sdk/templates/template_manager.rs
Normal file
@@ -0,0 +1,194 @@
|
||||
//! TemplateManager for managing workflow templates
|
||||
|
||||
use std::collections::HashMap;
|
||||
use crate::types::{WorkflowTemplateData, ParameterValues};
|
||||
use crate::templates::{WorkflowTemplate, WorkflowInstance};
|
||||
use crate::error::{ComfyUIError, Result};
|
||||
|
||||
/// Manages a collection of workflow templates
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct TemplateManager {
|
||||
templates: HashMap<String, WorkflowTemplate>,
|
||||
}
|
||||
|
||||
impl TemplateManager {
|
||||
/// Creates a new TemplateManager
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
templates: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Registers a template from template data
|
||||
pub fn register_from_data(&mut self, data: WorkflowTemplateData) -> Result<()> {
|
||||
let template = WorkflowTemplate::from_data(data)?;
|
||||
let id = template.id().to_string();
|
||||
self.templates.insert(id, template);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Registers a template
|
||||
pub fn register(&mut self, template: WorkflowTemplate) {
|
||||
let id = template.id().to_string();
|
||||
self.templates.insert(id, template);
|
||||
}
|
||||
|
||||
/// Gets a template by ID
|
||||
pub fn get_by_id(&self, id: &str) -> Option<&WorkflowTemplate> {
|
||||
self.templates.get(id)
|
||||
}
|
||||
|
||||
/// Gets a mutable reference to a template by ID
|
||||
pub fn get_by_id_mut(&mut self, id: &str) -> Option<&mut WorkflowTemplate> {
|
||||
self.templates.get_mut(id)
|
||||
}
|
||||
|
||||
/// Removes a template by ID
|
||||
pub fn remove(&mut self, id: &str) -> Option<WorkflowTemplate> {
|
||||
self.templates.remove(id)
|
||||
}
|
||||
|
||||
/// Lists all template IDs
|
||||
pub fn list_ids(&self) -> Vec<&str> {
|
||||
self.templates.keys().map(|s| s.as_str()).collect()
|
||||
}
|
||||
|
||||
/// Lists all templates
|
||||
pub fn list_templates(&self) -> Vec<&WorkflowTemplate> {
|
||||
self.templates.values().collect()
|
||||
}
|
||||
|
||||
/// Gets templates by category
|
||||
pub fn get_by_category(&self, category: &str) -> Vec<&WorkflowTemplate> {
|
||||
self.templates
|
||||
.values()
|
||||
.filter(|template| {
|
||||
template.category()
|
||||
.map(|c| c == category)
|
||||
.unwrap_or(false)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Gets templates by tag
|
||||
pub fn get_by_tag(&self, tag: &str) -> Vec<&WorkflowTemplate> {
|
||||
self.templates
|
||||
.values()
|
||||
.filter(|template| template.tags().contains(&tag.to_string()))
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Searches templates by name (case-insensitive partial match)
|
||||
pub fn search_by_name(&self, query: &str) -> Vec<&WorkflowTemplate> {
|
||||
let query_lower = query.to_lowercase();
|
||||
self.templates
|
||||
.values()
|
||||
.filter(|template| {
|
||||
template.name().to_lowercase().contains(&query_lower)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Searches templates by description (case-insensitive partial match)
|
||||
pub fn search_by_description(&self, query: &str) -> Vec<&WorkflowTemplate> {
|
||||
let query_lower = query.to_lowercase();
|
||||
self.templates
|
||||
.values()
|
||||
.filter(|template| {
|
||||
template.description()
|
||||
.map(|desc| desc.to_lowercase().contains(&query_lower))
|
||||
.unwrap_or(false)
|
||||
})
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Creates a workflow instance from a template
|
||||
pub fn create_instance(&self, template_id: &str, parameters: ParameterValues) -> Result<WorkflowInstance> {
|
||||
let template = self.get_by_id(template_id)
|
||||
.ok_or_else(|| ComfyUIError::template_validation(
|
||||
format!("Template with ID '{}' not found", template_id)
|
||||
))?;
|
||||
|
||||
template.create_instance(parameters)
|
||||
}
|
||||
|
||||
/// Gets the number of registered templates
|
||||
pub fn count(&self) -> usize {
|
||||
self.templates.len()
|
||||
}
|
||||
|
||||
/// Checks if a template exists
|
||||
pub fn contains(&self, id: &str) -> bool {
|
||||
self.templates.contains_key(id)
|
||||
}
|
||||
|
||||
/// Clears all templates
|
||||
pub fn clear(&mut self) {
|
||||
self.templates.clear();
|
||||
}
|
||||
|
||||
/// Gets all categories
|
||||
pub fn get_categories(&self) -> Vec<String> {
|
||||
let mut categories: Vec<String> = self.templates
|
||||
.values()
|
||||
.filter_map(|template| template.category().map(|c| c.to_string()))
|
||||
.collect();
|
||||
categories.sort();
|
||||
categories.dedup();
|
||||
categories
|
||||
}
|
||||
|
||||
/// Gets all tags
|
||||
pub fn get_tags(&self) -> Vec<String> {
|
||||
let mut tags: Vec<String> = self.templates
|
||||
.values()
|
||||
.flat_map(|template| template.tags().iter().cloned())
|
||||
.collect();
|
||||
tags.sort();
|
||||
tags.dedup();
|
||||
tags
|
||||
}
|
||||
|
||||
/// Validates all templates
|
||||
pub fn validate_all(&self) -> Result<()> {
|
||||
for (id, template) in &self.templates {
|
||||
// Try to create an instance with default values to validate template
|
||||
let default_values = template.get_default_values();
|
||||
let validation = template.validate(&default_values);
|
||||
|
||||
if !validation.valid {
|
||||
return Err(ComfyUIError::template_validation(
|
||||
format!("Template '{}' validation failed", id)
|
||||
));
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Exports all templates as template data
|
||||
pub fn export_all(&self) -> Vec<WorkflowTemplateData> {
|
||||
self.templates
|
||||
.values()
|
||||
.map(|template| template.to_data())
|
||||
.collect()
|
||||
}
|
||||
|
||||
/// Imports templates from template data
|
||||
pub fn import_all(&mut self, templates: Vec<WorkflowTemplateData>) -> Result<Vec<String>> {
|
||||
let mut imported_ids = Vec::new();
|
||||
|
||||
for data in templates {
|
||||
let id = data.metadata.id.clone();
|
||||
self.register_from_data(data)?;
|
||||
imported_ids.push(id);
|
||||
}
|
||||
|
||||
Ok(imported_ids)
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for TemplateManager {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user