1
1
using System ;
2
2
using System . Collections . Generic ;
3
3
using System . Collections . Immutable ;
4
+ using System . ComponentModel . DataAnnotations ;
4
5
using System . Diagnostics ;
5
6
using System . Diagnostics . CodeAnalysis ;
6
7
using System . IO ;
7
8
using System . Linq ;
8
- using System . Management ;
9
9
using System . Text . Json ;
10
10
using System . Text . Json . Serialization ;
11
11
using System . Threading ;
15
15
using Avalonia . Threading ;
16
16
using CommunityToolkit . Mvvm . Input ;
17
17
using ExifLibrary ;
18
- using MetadataExtractor . Formats . Exif ;
19
18
using NLog ;
20
19
using Refit ;
21
20
using SkiaSharp ;
27
26
using StabilityMatrix . Avalonia . ViewModels . Dialogs ;
28
27
using StabilityMatrix . Avalonia . ViewModels . Inference ;
29
28
using StabilityMatrix . Avalonia . ViewModels . Inference . Modules ;
30
- using StabilityMatrix . Core . Animation ;
31
29
using StabilityMatrix . Core . Exceptions ;
32
30
using StabilityMatrix . Core . Extensions ;
33
31
using StabilityMatrix . Core . Helper ;
@@ -297,14 +295,18 @@ protected async Task RunGeneration(ImageGenerationEventArgs args, CancellationTo
297
295
Task . Run (
298
296
async ( ) =>
299
297
{
300
- var delayTime = 250 - ( int ) timer . ElapsedMilliseconds ;
301
- if ( delayTime > 0 )
298
+ try
302
299
{
303
- await Task . Delay ( delayTime , cancellationToken ) ;
300
+ var delayTime = 250 - ( int ) timer . ElapsedMilliseconds ;
301
+ if ( delayTime > 0 )
302
+ {
303
+ await Task . Delay ( delayTime , cancellationToken ) ;
304
+ }
305
+
306
+ // ReSharper disable once AccessToDisposedClosure
307
+ AttachRunningNodeChangedHandler ( promptTask ) ;
304
308
}
305
-
306
- // ReSharper disable once AccessToDisposedClosure
307
- AttachRunningNodeChangedHandler ( promptTask ) ;
309
+ catch ( TaskCanceledException ) { }
308
310
} ,
309
311
cancellationToken
310
312
)
@@ -328,10 +330,7 @@ await DialogHelper
328
330
// Get output images
329
331
var imageOutputs = await client . GetImagesForExecutedPromptAsync ( promptTask . Id , cancellationToken ) ;
330
332
331
- if (
332
- ! imageOutputs . TryGetValue ( args . OutputNodeNames [ 0 ] , out var images )
333
- || images is not { Count : > 0 }
334
- )
333
+ if ( imageOutputs . Values . All ( images => images is null or { Count : 0 } ) )
335
334
{
336
335
// No images match
337
336
notificationService . Show (
@@ -350,7 +349,7 @@ await DialogHelper
350
349
ImageGalleryCardViewModel . ImageSources . Clear ( ) ;
351
350
}
352
351
353
- var outputImages = await ProcessOutputImages ( images , args ) ;
352
+ var outputImages = await ProcessAllOutputImages ( imageOutputs , args ) ;
354
353
355
354
var notificationImage = outputImages . FirstOrDefault ( ) ? . LocalFile ;
356
355
@@ -380,12 +379,34 @@ await notificationService.ShowAsync(
380
379
}
381
380
}
382
381
382
+ private async Task < IEnumerable < ImageSource > > ProcessAllOutputImages (
383
+ IReadOnlyDictionary < string , List < ComfyImage > ? > images ,
384
+ ImageGenerationEventArgs args
385
+ )
386
+ {
387
+ var results = new List < ImageSource > ( ) ;
388
+
389
+ foreach ( var ( nodeName , imageList ) in images )
390
+ {
391
+ if ( imageList is null )
392
+ {
393
+ Logger . Warn ( "No images for node {NodeName}" , nodeName ) ;
394
+ continue ;
395
+ }
396
+
397
+ results . AddRange ( await ProcessOutputImages ( imageList , args , nodeName . Replace ( '_' , ' ' ) ) ) ;
398
+ }
399
+
400
+ return results ;
401
+ }
402
+
383
403
/// <summary>
384
404
/// Handles image output metadata for generation runs
385
405
/// </summary>
386
406
private async Task < List < ImageSource > > ProcessOutputImages (
387
407
IReadOnlyCollection < ComfyImage > images ,
388
- ImageGenerationEventArgs args
408
+ ImageGenerationEventArgs args ,
409
+ string ? imageLabel = null
389
410
)
390
411
{
391
412
var client = args . Client ;
@@ -441,7 +462,7 @@ ImageGenerationEventArgs args
441
462
images . Count
442
463
) ;
443
464
444
- outputImages . Add ( new ImageSource ( filePath ) ) ;
465
+ outputImages . Add ( new ImageSource ( filePath ) { Label = imageLabel } ) ;
445
466
EventManager . Instance . OnImageFileAdded ( filePath ) ;
446
467
}
447
468
else if ( comfyImage . FileName . EndsWith ( ".webp" ) )
@@ -470,7 +491,7 @@ ImageGenerationEventArgs args
470
491
fileExtension : Path . GetExtension ( comfyImage . FileName ) . Replace ( "." , "" )
471
492
) ;
472
493
473
- outputImages . Add ( new ImageSource ( filePath ) ) ;
494
+ outputImages . Add ( new ImageSource ( filePath ) { Label = imageLabel } ) ;
474
495
EventManager . Instance . OnImageFileAdded ( filePath ) ;
475
496
}
476
497
else
@@ -484,7 +505,7 @@ ImageGenerationEventArgs args
484
505
fileExtension : Path . GetExtension ( comfyImage . FileName ) . Replace ( "." , "" )
485
506
) ;
486
507
487
- outputImages . Add ( new ImageSource ( filePath ) ) ;
508
+ outputImages . Add ( new ImageSource ( filePath ) { Label = imageLabel } ) ;
488
509
EventManager . Instance . OnImageFileAdded ( filePath ) ;
489
510
}
490
511
}
@@ -554,7 +575,12 @@ private async Task GenerateImage(
554
575
}
555
576
catch ( OperationCanceledException )
556
577
{
557
- Logger . Debug ( $ "Image Generation Canceled") ;
578
+ Logger . Debug ( "Image Generation Canceled" ) ;
579
+ }
580
+ catch ( ValidationException e )
581
+ {
582
+ Logger . Debug ( "Image Generation Validation Error: {Message}" , e . Message ) ;
583
+ notificationService . Show ( "Validation Error" , e . Message , NotificationType . Error ) ;
558
584
}
559
585
}
560
586
0 commit comments