diff --git a/StabilityMatrix.Core/Models/Api/Comfy/NodeTypes/NodeConnections.cs b/StabilityMatrix.Core/Models/Api/Comfy/NodeTypes/NodeConnections.cs index d5630c9f..ce6521fe 100644 --- a/StabilityMatrix.Core/Models/Api/Comfy/NodeTypes/NodeConnections.cs +++ b/StabilityMatrix.Core/Models/Api/Comfy/NodeTypes/NodeConnections.cs @@ -15,3 +15,5 @@ public class ModelNodeConnection : NodeConnectionBase { } public class ConditioningNodeConnection : NodeConnectionBase { } public class ClipNodeConnection : NodeConnectionBase { } + +public class ControlNetNodeConnection : NodeConnectionBase { } diff --git a/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs b/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs index 0cbf37fe..1cf1947d 100644 --- a/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs +++ b/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyNodeBuilder.cs @@ -348,6 +348,18 @@ public class ComfyNodeBuilder }; } + public class ControlNetApplyAdvanced + : ComfyTypedNodeBase + { + public required ConditioningNodeConnection Positive { get; init; } + public required ConditioningNodeConnection Negative { get; init; } + public required ControlNetNodeConnection ControlNet { get; init; } + public required ImageNodeConnection Image { get; init; } + public required double Strength { get; init; } + public required double StartPercent { get; init; } + public required double EndPercent { get; init; } + } + public ImageNodeConnection Lambda_LatentToImage( LatentNodeConnection latent, VAENodeConnection vae @@ -796,21 +808,10 @@ public class ComfyNodeBuilder public VAENodeConnection? PrimaryVAE { get; set; } public Size PrimarySize { get; set; } - /*public LatentNodeConnection? Latent { get; set; } - public Size LatentSize { get; set; } - - public ImageNodeConnection? Image { get; set; } - public Size ImageSize { get; set; }*/ - public List OutputNodes { get; } = new(); public IEnumerable OutputNodeNames => OutputNodes.Select(n => n.Name); - public VAENodeConnection GetRefinerOrBaseVAE() - { - return RefinerVAE ?? BaseVAE ?? throw new NullReferenceException("No VAE"); - } - public ModelNodeConnection GetRefinerOrBaseModel() { return RefinerModel ?? BaseModel ?? throw new NullReferenceException("No Model"); diff --git a/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyTypedNodeBase.cs b/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyTypedNodeBase.cs new file mode 100644 index 00000000..ad5a3145 --- /dev/null +++ b/StabilityMatrix.Core/Models/Api/Comfy/Nodes/ComfyTypedNodeBase.cs @@ -0,0 +1,65 @@ +using System.Reflection; +using System.Text.Json.Serialization; +using StabilityMatrix.Core.Models.Api.Comfy.NodeTypes; +using Yoh.Text.Json.NamingPolicies; + +namespace StabilityMatrix.Core.Models.Api.Comfy.Nodes; + +public abstract class ComfyTypedNodeBase +{ + protected virtual string ClassType => GetType().Name; + + [JsonIgnore] + public required string Name { get; init; } + + protected NamedComfyNode ToNamedNode() + { + var inputs = new Dictionary(); + + // Loop through all properties, key is property name as snake_case, or JsonPropertyName + var namingPolicy = JsonNamingPolicies.SnakeCaseLower; + + foreach (var property in GetType().GetProperties()) + { + if (property.Name == nameof(Name) || property.GetValue(this) is not { } value) + continue; + + // Skip JsonIgnore + if (property.GetCustomAttribute() is not null) + continue; + + var key = + property.GetCustomAttribute()?.Name + ?? namingPolicy.ConvertName(property.Name); + + inputs.Add(key, value); + } + + return new NamedComfyNode(Name) { ClassType = ClassType, Inputs = inputs }; + } + + // Implicit conversion to NamedComfyNode + public static implicit operator NamedComfyNode(ComfyTypedNodeBase node) => node.ToNamedNode(); +} + +public abstract class ComfyTypedNodeBase : ComfyTypedNodeBase + where TOutput : NodeConnectionBase, new() +{ + public TOutput Output => new() { Data = new object[] { Name, 0 } }; + + public static implicit operator NamedComfyNode(ComfyTypedNodeBase node) => + (NamedComfyNode)node.ToNamedNode(); +} + +public abstract class ComfyTypedNodeBase : ComfyTypedNodeBase + where TOutput1 : NodeConnectionBase, new() + where TOutput2 : NodeConnectionBase, new() +{ + public TOutput1 Output1 => new() { Data = new object[] { Name, 0 } }; + + public TOutput1 Output2 => new() { Data = new object[] { Name, 1 } }; + + public static implicit operator NamedComfyNode( + ComfyTypedNodeBase node + ) => (NamedComfyNode)node.ToNamedNode(); +}