Skip to content

Commit 9d8e881

Browse files
mythzmythz
committed
refactor media models
1 parent 6c7f131 commit 9d8e881

File tree

18 files changed

+1375
-556
lines changed

18 files changed

+1375
-556
lines changed

‎AiServer.ServiceInterface/AppData.cs‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ public void ResetMediaProviders(IDbConnection db)
119119
private void LoadModelDefaults()
120120
{
121121
MediaModels = LoadModels<MediaModel>("media-models.json")
122-
.Where(x => x is { ApiModels.Keys.Count: > 0, ModelType:
122+
.Where(x => x is { ApiModels.Keys.Count: > 0, Type:
123123
ModelType.TextToImage or
124124
ModelType.TextToSpeech or
125125
ModelType.SpeechToText or
@@ -159,7 +159,7 @@ public MediaModel GetMediaModelByApiModel(MediaProvider provider, string apiMode
159159
public string? GetQualifiedMediaModel(ModelType modelType, string apiModel)
160160
{
161161
foreach (var mediaModel in MediaModels
162-
.Where(x => x.ModelType == modelType))
162+
.Where(x => x.Type == modelType))
163163
{
164164
foreach (var entry in mediaModel.ApiModels)
165165
{
@@ -174,7 +174,7 @@ public MediaModel GetMediaModelByApiModel(MediaProvider provider, string apiMode
174174
{
175175
ArgumentNullException.ThrowIfNull(provider);
176176
var supportedTaskModels = MediaModelsMap.Values
177-
.Where(x => x.ModelType == GetModelTypeByAiTaskType(taskType))
177+
.Where(x => x.Type == GetModelTypeByAiTaskType(taskType))
178178
.Where(x => x.ApiModels.ContainsKey(provider.MediaType!.Id) &&
179179
provider.Models != null &&
180180
provider.Models.Contains(x.ApiModels[provider.MediaType!.Id]));
@@ -187,7 +187,7 @@ public MediaModel GetMediaModelByApiModel(MediaProvider provider, string apiMode
187187
public List<string> GetSupportedModels(AiTaskType taskType)
188188
{
189189
return MediaModelsMap
190-
.Where(x => x.Value.ModelType == GetModelTypeByAiTaskType(taskType))
190+
.Where(x => x.Value.Type == GetModelTypeByAiTaskType(taskType))
191191
.Select(x => x.Key)
192192
.ToList();
193193
}
@@ -196,7 +196,7 @@ public bool ProviderHasModelForTask(MediaProvider provider, AiTaskType taskType)
196196
{
197197
ArgumentNullException.ThrowIfNull(provider);
198198
return MediaModelsMap.Values
199-
.Any(x => x.ModelType == GetModelTypeByAiTaskType(taskType) &&
199+
.Any(x => x.Type == GetModelTypeByAiTaskType(taskType) &&
200200
x.ApiModels.ContainsKey(provider.MediaType!.Id) &&
201201
provider.Models != null &&
202202
provider.Models.Contains(x.ApiModels[provider.MediaType!.Id]));
@@ -205,7 +205,7 @@ public bool ProviderHasModelForTask(MediaProvider provider, AiTaskType taskType)
205205
public bool ModelSupportsTask(string modelId, AiTaskType taskType)
206206
{
207207
return MediaModelsMap.TryGetValue(modelId, out var modelSettings) &&
208-
modelSettings.ModelType == GetModelTypeByAiTaskType(taskType);
208+
modelSettings.Type == GetModelTypeByAiTaskType(taskType);
209209
}
210210

211211
private ModelType GetModelTypeByAiTaskType(AiTaskType taskType)

‎AiServer.ServiceInterface/Comfy/ComfyClient.cs‎

Lines changed: 21 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -127,7 +127,7 @@ public async Task<ComfyFileInput> UploadAudioAssetAsync(Stream fileStream, strin
127127
return result.FromJson<ComfyFileInput>();
128128
}
129129

130-
public async Task<string> PopulateWorkflowAsync(ComfyWorkflowRequest request, string templatePath, CancellationToken token)
130+
public async Task<string> PopulateWorkflowAsync(Dictionary<string,object> args, string templatePath, CancellationToken token)
131131
{
132132
var buildInTemplatePath = Path.Combine(WorkflowTemplatePath, templatePath);
133133
var overrideTemplatePath = Path.Combine("App_Data","overrides", templatePath);
@@ -138,10 +138,10 @@ public async Task<string> PopulateWorkflowAsync(ComfyWorkflowRequest request, st
138138
throw new Exception($"Template file not found: {useTemplatePath}");
139139
// Read template from file for Text to Image
140140
var template = await File.ReadAllTextAsync(useTemplatePath, token);
141+
141142
// Populate template with request
142-
var workflowPageResult = new PageResult(context.OneTimePage(template))
143-
{
144-
Args = request.ToObjectDictionary(),
143+
var workflowPageResult = new PageResult(context.OneTimePage(template)) {
144+
Args = args,
145145
};
146146

147147
// Render template to JSON
@@ -181,7 +181,23 @@ public async Task<ComfyWorkflowResponse> PromptGenerationAsync(ComfyWorkflowRequ
181181
await HandleAssetUploadsAsync(comfyRequest, promptId, token);
182182

183183
// Read template from file for Text to Image
184-
var workflowJson = await PopulateWorkflowAsync(comfyRequest, templatePath, token);
184+
var workflowArgs = comfyRequest.ToObjectDictionary();
185+
foreach (var key in workflowArgs.Keys.ToList())
186+
{
187+
workflowArgs[key.ToCamelCase()] = workflowArgs[key];
188+
}
189+
190+
var mediaModel = AppData.Instance.MediaModels.FirstOrDefault(x =>
191+
x.ApiModels?.TryGetValue("ComfyUI", out var comfyModel) == true && comfyModel == comfyRequest.Model);
192+
if (mediaModel?.WorkflowVars != null)
193+
{
194+
foreach (var entry in mediaModel.WorkflowVars)
195+
{
196+
workflowArgs[entry.Key] = entry.Value;
197+
}
198+
}
199+
200+
var workflowJson = await PopulateWorkflowAsync(workflowArgs, templatePath, token);
185201
// Convert to ComfyUI API JSON format
186202
var apiJson = await ConvertWorkflowToApiAsync(workflowJson, token);
187203

‎AiServer.ServiceInterface/GenerationServices.cs‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -77,13 +77,14 @@ public object Any(ActiveMediaModels request)
7777
.ThenBy(x => x.Name)
7878
.SelectMany(x =>
7979
x.Models.Select(m => appData.GetQualifiedMediaModel(ModelType.TextToImage, m)))
80+
.Select(x => x == null ? null : appData.MediaModels.FirstOrDefault(m => m.Id == x))
8081
.Where(x => x != null)
81-
.Select(x => x!) // Non-null assertion after filtering out null values
82+
.Select(x => x!)
8283
.Distinct();
83-
84-
return new StringsResponse
84+
85+
return new ActiveMediaModelsResponse
8586
{
86-
Results = activeModels.ToList()
87+
Results = activeModels.Map(x => new Entry { Key = x.Id, Value = x.Name ?? x.Id })
8788
};
8889
}
8990

‎AiServer.ServiceModel/ComfyPublic.cs‎

Lines changed: 0 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -154,27 +154,6 @@ namespace AiServer.ServiceModel;
154154
// public List<ComfyHostedFileOutput>? Sounds { get; set; } = new();
155155
// }
156156

157-
[Tag(Tags.Media)]
158-
public class ConfigureAndDownloadModel : IReturn<ComfyAgentDownloadStatus>
159-
{
160-
public string Name { get; set; }
161-
public string Filename { get; set; }
162-
public string DownloadUrl { get; set; }
163-
164-
public double? CfgScale { get; set; }
165-
166-
public string? Scheduler { get; set; }
167-
168-
public ComfySampler? Sampler { get; set; }
169-
170-
public int? Width { get; set; }
171-
172-
public int? Height { get; set; }
173-
174-
public int? Steps { get; set; }
175-
176-
public string? NegativePrompt { get; set; }
177-
}
178157
public class ComfyAgentDownloadStatus
179158
{
180159
public string? Name { get; set; }

‎AiServer.ServiceModel/Generations.cs‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,18 @@ namespace AiServer.ServiceModel;
55

66
[Tag(Tags.AiInfo)]
77
[Api("Active Media Worker Models available in AI Server")]
8-
public class ActiveMediaModels : IGet, IReturn<StringsResponse> {}
8+
public class ActiveMediaModels : IGet, IReturn<ActiveMediaModelsResponse> {}
9+
public class ActiveMediaModelsResponse
10+
{
11+
public List<Entry> Results { get; set; } = [];
12+
public ResponseStatus? ResponseStatus { get; set; }
13+
}
14+
15+
public class Entry
16+
{
17+
public string Key { get; set; }
18+
public string Value { get; set; }
19+
}
920

1021

1122
[ValidateApiKey]

‎AiServer.ServiceModel/Types/Generations.cs‎

Lines changed: 12 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -132,26 +132,18 @@ public enum AiServiceProvider
132132
public class MediaModel : IHasId<string>
133133
{
134134
public string Id { get; set; }
135+
public string? Name { get; set; }
136+
public ModelType? Type { get; set; }
135137
public Dictionary<string, string> ApiModels { get; set; } = new();
136-
public string? Url { get; set; }
137-
public double? Quality { get; set; }
138-
public string? AspectRatio { get; set; }
139-
public double? CfgScale { get; set; }
140-
public string? Scheduler { get; set; }
141-
public ComfySampler? Sampler { get; set; }
142-
public int? Width { get; set; }
143-
public int? Height { get; set; }
144-
public int? Steps { get; set; }
145-
public string? NegativePrompt { get; set; }
146-
public ModelType? ModelType { get; set; }
147-
public string? Filename { get; set; }
148-
public string? Workflow { get; set; }
149-
/// <summary>
150-
/// Flag for models that are pulled down when run for the first time
151-
/// Keys are the ProviderId strings
152-
/// </summary>
153-
public Dictionary<string,bool>? OnDemand { get; set; }
154138
public Dictionary<string,List<string>>? SupportedTasks { get; set; }
139+
public List<string>? Dependencies { get; set; }
140+
public string? Installer { get; set; }
141+
public string? Path { get; set; }
142+
public string? Workflow { get; set; }
143+
public Dictionary<string, object> WorkflowVars { get; set; } = new();
144+
public string? DownloadToken { get; set; }
145+
public string? DownloadUrl { get; set; }
146+
public string? Url { get; set; }
155147
}
156148

157149
public enum ModelType
@@ -165,7 +157,8 @@ public enum ModelType
165157
ImageToText,
166158
ImageToImage,
167159
ImageWithMask,
168-
VAE
160+
Lora,
161+
VAE,
169162
}
170163

171164
public class TextToSpeechVoice : IHasId<string>

‎AiServer.Tests/ComfyUITests.cs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -644,7 +644,7 @@ public async Task Can_use_ImageToImageUpscale_from_template_workflow()
644644
};
645645

646646
// Convert template to JSON
647-
var jsonTemplate = await client.PopulateWorkflowAsync(testDto,client.DefaultImageToImageUpscaleTemplate,cts.Token);
647+
var jsonTemplate = await client.PopulateWorkflowAsync(testDto.ToObjectDictionary(),client.DefaultImageToImageUpscaleTemplate,cts.Token);
648648

649649
// Assert that the JSON template is not null
650650
Assert.That(jsonTemplate, Is.Not.Null);
@@ -684,7 +684,7 @@ public async Task Can_use_ImageToImage_from_template_workflow()
684684
};
685685

686686
// Convert template to JSON
687-
var jsonTemplate = await client.PopulateWorkflowAsync(testDto,client.DefaultImageToImageTemplate, cts.Token);
687+
var jsonTemplate = await client.PopulateWorkflowAsync(testDto.ToObjectDictionary(),client.DefaultImageToImageTemplate, cts.Token);
688688

689689
// Assert that the JSON template is not null
690690
Assert.That(jsonTemplate, Is.Not.Null);
@@ -721,7 +721,7 @@ public async Task Can_use_TextToImage_from_template_workflow()
721721
};
722722

723723
// Convert template to JSON
724-
var jsonTemplate = await client.PopulateWorkflowAsync(testDto,client.DefaultTextToImageTemplate, cts.Token);
724+
var jsonTemplate = await client.PopulateWorkflowAsync(testDto.ToObjectDictionary(),client.DefaultTextToImageTemplate, cts.Token);
725725

726726
// Assert that the JSON template is not null
727727
Assert.That(jsonTemplate, Is.Not.Null);

‎AiServer/Configure.BackgroundJobs.cs‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,9 +34,12 @@ public void Configure(IWebHostBuilder builder) => builder
3434
AppData.Instance.Reload(db);
3535

3636
var mediaOptions = services.GetRequiredService<ComfyMediaProviderOptions>();
37-
foreach (var model in AppData.Instance.MediaModels.Where(x => x.Filename != null && x.Workflow != null))
37+
foreach (var model in AppData.Instance.MediaModels.Where(x => x.Path != null && x.Workflow != null))
3838
{
39-
mediaOptions.TextToImageModelOverrides[model.Filename!] = model.Workflow!;
39+
if (model.ApiModels.TryGetValue("ComfyUI", out var comfyModel))
40+
{
41+
mediaOptions.TextToImageModelOverrides[comfyModel] = model.Workflow!;
42+
}
4043
}
4144

4245
var jobs = services.GetRequiredService<IBackgroundJobs>();

AiServer/workflows/flux1/text_to_image.json renamed to AiServer/workflows/text-to-image/flux1.json

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -74,9 +74,9 @@
7474
"Node name for S&R": "EmptyLatentImage"
7575
},
7676
"widgets_values": [
77-
1024,
78-
1024,
79-
1
77+
{{width}},
78+
{{height}},
79+
{{batchSize}}
8080
],
8181
"color": "#323",
8282
"bgcolor": "#535"
@@ -109,7 +109,7 @@
109109
"Node name for S&R": "RandomNoise"
110110
},
111111
"widgets_values": [
112-
{{Seed}},
112+
{{seed}},
113113
"randomize"
114114
],
115115
"color": "#2a363b",
@@ -143,7 +143,7 @@
143143
"Node name for S&R": "KSamplerSelect"
144144
},
145145
"widgets_values": [
146-
"{{Sampler}}"
146+
"{{sampler}}"
147147
]
148148
},
149149
{
@@ -180,8 +180,8 @@
180180
"Node name for S&R": "BasicScheduler"
181181
},
182182
"widgets_values": [
183-
"{{Scheduler}}",
184-
{{Steps}},
183+
"{{scheduler}}",
184+
{{steps}},
185185
1
186186
]
187187
},
@@ -356,7 +356,7 @@
356356
"Node name for S&R": "CLIPTextEncode"
357357
},
358358
"widgets_values": [
359-
"{{PositivePrompt}}"
359+
"{{positivePrompt}}"
360360
],
361361
"color": "#232",
362362
"bgcolor": "#353"

AiServer/workflows/hidream/text_to_image.json renamed to AiServer/workflows/text-to-image/hidream.json

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -42,9 +42,9 @@
4242
"secondTabWidth": 65
4343
},
4444
"widgets_values": [
45-
{{Width}},
46-
{{Height}},
47-
{{BatchSize}}
45+
{{width}},
46+
{{height}},
47+
{{batchSize}}
4848
]
4949
},
5050
{
@@ -144,7 +144,7 @@
144144
"secondTabWidth": 65
145145
},
146146
"widgets_values": [
147-
"{{NegativePrompt}}"
147+
"{{negativePrompt}}"
148148
],
149149
"color": "#322",
150150
"bgcolor": "#533"
@@ -292,12 +292,12 @@
292292
"secondTabWidth": 65
293293
},
294294
"widgets_values": [
295-
{{Seed}},
295+
{{seed}},
296296
"randomize",
297-
{{Steps}},
297+
{{steps}},
298298
{{CfgScale}},
299-
"{{Sampler}}",
300-
"{{Scheduler}}",
299+
"{{sampler}}",
300+
"{{scheduler}}",
301301
1
302302
]
303303
},
@@ -505,7 +505,7 @@
505505
"secondTabWidth": 65
506506
},
507507
"widgets_values": [
508-
"{{PositivePrompt}}"
508+
"{{positivePrompt}}"
509509
],
510510
"color": "#232",
511511
"bgcolor": "#353"

0 commit comments

Comments
 (0)