From 28a0014fa0c87853e315e91c3012ed4914441db1 Mon Sep 17 00:00:00 2001
From: Menno van Leeuwen <menno@vleeuwen.me>
Date: Sun, 30 Mar 2025 18:40:31 +0200
Subject: [PATCH] adds proper support for fetching images

---
 lib/src/_http_endpoints.dart               | 17 +++++++++++++----
 lib/src/comfyui_api.dart                   |  7 ++++---
 lib/src/models/execution_event.dart        | 10 ++++++++++
 lib/src/utils/websocket_event_handler.dart | 17 +++++++++++++++++
 4 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/lib/src/_http_endpoints.dart b/lib/src/_http_endpoints.dart
index ef8b84a..7e20e94 100644
--- a/lib/src/_http_endpoints.dart
+++ b/lib/src/_http_endpoints.dart
@@ -42,10 +42,19 @@ class ComfyApiHttpEndpoints {
     return HistoryResponse.fromJson(jsonDecode(response.body));
   }
 
-  /// Gets image data by filename
-  Future<List<int>> getImage(String filename) async {
-    final response =
-        await _httpClient.get(Uri.parse('$host/api/view?filename=$filename'));
+  /// Gets image data by filename, optionally specifying subfolder and type
+  Future<List<int>> getImage(String filename,
+      {String? subfolder, String? type}) async {
+    // Build query parameters dynamically
+    final queryParameters = {
+      'filename': filename,
+      if (subfolder != null && subfolder.isNotEmpty) 'subfolder': subfolder,
+      if (type != null && type.isNotEmpty) 'type': type,
+    };
+    final uri =
+        Uri.parse('$host/api/view').replace(queryParameters: queryParameters);
+    logger.d('Fetching image from URI: $uri'); // Log the final URI
+    final response = await _httpClient.get(uri);
     _validateResponse(response);
     return response.bodyBytes;
   }
diff --git a/lib/src/comfyui_api.dart b/lib/src/comfyui_api.dart
index b4edcfa..251eba1 100644
--- a/lib/src/comfyui_api.dart
+++ b/lib/src/comfyui_api.dart
@@ -96,9 +96,10 @@ class ComfyUiApi {
   Future<HistoryResponse> getHistory({int maxItems = 64}) =>
       _httpEndpoints.getHistory(maxItems: maxItems);
 
-  /// Gets image data by filename
-  Future<List<int>> getImage(String filename) =>
-      _httpEndpoints.getImage(filename);
+  /// Gets image data by filename, optionally specifying subfolder and type
+  Future<List<int>> getImage(String filename,
+          {String? subfolder, String? type}) =>
+      _httpEndpoints.getImage(filename, subfolder: subfolder, type: type);
 
   /// Gets a list of all available models
   Future<Map<String, dynamic>> getModels() => _httpEndpoints.getModels();
diff --git a/lib/src/models/execution_event.dart b/lib/src/models/execution_event.dart
index 909f402..ef2a6fd 100644
--- a/lib/src/models/execution_event.dart
+++ b/lib/src/models/execution_event.dart
@@ -2,20 +2,29 @@ class ExecutionEvent {
   final String promptId;
   final int timestamp;
   final String? node;
+  final List<Map<String, dynamic>>? outputImages;
   final Map<String, dynamic>? extra;
 
   const ExecutionEvent({
     required this.promptId,
     required this.timestamp,
     this.node,
+    this.outputImages,
     this.extra,
   });
 
   factory ExecutionEvent.fromJson(Map<String, dynamic> json) {
+    // Note: This factory is not directly used by the handler,
+    // but kept for potential future use or consistency.
+    // The handler now constructs the event directly.
     return ExecutionEvent(
       promptId: json['prompt_id'] as String,
       timestamp: json['timestamp'] as int,
       node: json['node'] as String?,
+      // Parsing logic moved to the handler for direct use.
+      // outputImages: (json['output']?['images'] as List<dynamic>?)
+      //     ?.map((e) => e as Map<String, dynamic>)
+      //     .toList(),
       extra: json['extra'] as Map<String, dynamic>?,
     );
   }
@@ -25,6 +34,7 @@ class ExecutionEvent {
       'prompt_id': promptId,
       'timestamp': timestamp,
       'node': node,
+      'outputImages': outputImages, // Added for completeness
       'extra': extra,
     };
   }
diff --git a/lib/src/utils/websocket_event_handler.dart b/lib/src/utils/websocket_event_handler.dart
index e261765..e1e6bf9 100644
--- a/lib/src/utils/websocket_event_handler.dart
+++ b/lib/src/utils/websocket_event_handler.dart
@@ -38,11 +38,28 @@ class WebSocketEventHandler {
   ) {
     if (event.data.containsKey('prompt_id')) {
       try {
+        List<Map<String, dynamic>>? outputImages;
+        if (event.data.containsKey('output') &&
+            event.data['output'] is Map &&
+            (event.data['output'] as Map).containsKey('images') &&
+            event.data['output']['images'] is List) {
+          // Ensure correct casting
+          try {
+            outputImages = (event.data['output']['images'] as List<dynamic>)
+                .map((e) => Map<String, dynamic>.from(e as Map))
+                .toList();
+          } catch (e) {
+            print('Error parsing output images: $e');
+            outputImages = null; // Set to null if parsing fails
+          }
+        }
+
         final executionEvent = ExecutionEvent(
           promptId: event.data['prompt_id'] as String,
           timestamp: event.data['timestamp'] as int? ??
               DateTime.now().millisecondsSinceEpoch,
           node: event.data['node']?.toString(),
+          outputImages: outputImages, // Pass the extracted images
           extra: event.data['extra'] as Map<String, dynamic>?,
         );
         executionEventController.add(executionEvent);