comfyui_api_sdk/comfyui-api-sdk.dart

import 'dart:async';
import 'dart:convert';

import 'package:http/http.dart' as http;
import 'package:uuid/uuid.dart';
import 'package:web_socket_channel/web_socket_channel.dart';

class ComfyUiApi {
  final String host;
  final String clientId;
  final http.Client _httpClient;
  WebSocketChannel? _wsChannel;
  final StreamController<Map<String, dynamic>> _progressController =
      StreamController.broadcast();

  /// Stream of progress updates from ComfyUI
  Stream<Map<String, dynamic>> get progressUpdates =>
      _progressController.stream;

  /// Creates a new ComfyUI API client
  ///
  /// [host] The host of the ComfyUI server (e.g. 'http://localhost:8188')
  /// [clientId] Optional client ID, will be automatically generated if not provided
  ComfyUiApi({
    required this.host,
    String? clientId,
    http.Client? httpClient,
  })  : clientId = clientId ?? const Uuid().v4(),
        _httpClient = httpClient ?? http.Client();

  /// Connects to the WebSocket for progress updates
  Future<void> connectWebSocket() async {
    final wsUrl =
        'ws://${host.replaceFirst(RegExp(r'^https?://'), '')}/ws?clientId=$clientId';
    _wsChannel = WebSocketChannel.connect(Uri.parse(wsUrl));

    _wsChannel!.stream.listen((message) {
      final data = jsonDecode(message);
      _progressController.add(data);
    }, onError: (error) {
      print('WebSocket error: $error');
    }, onDone: () {
      print('WebSocket connection closed');
    });
  }

  /// Closes the WebSocket connection and cleans up resources
  void dispose() {
    _wsChannel?.sink.close();
    _progressController.close();
    _httpClient.close();
  }

  /// Gets the current queue status
  Future<Map<String, dynamic>> getQueue() async {
    final response = await _httpClient.get(Uri.parse('$host/queue'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets the history of the queue
  Future<Map<String, dynamic>> getHistory({int maxItems = 64}) async {
    final response = await _httpClient
        .get(Uri.parse('$host/api/history?max_items=$maxItems'));
    _validateResponse(response);
    return 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'));
    _validateResponse(response);
    return response.bodyBytes;
  }

  /// Gets a list of all available models
  Future<Map<String, dynamic>> getModels() async {
    final response =
        await _httpClient.get(Uri.parse('$host/api/experiment/models'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets a list of checkpoints
  Future<Map<String, dynamic>> getCheckpoints() async {
    final response = await _httpClient
        .get(Uri.parse('$host/api/experiment/models/checkpoints'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets details for a specific checkpoint
  Future<Map<String, dynamic>> getCheckpointDetails(
      String pathAndFileName) async {
    final response = await _httpClient.get(Uri.parse(
        '$host/api/view_metadata/checkpoints?filename=$pathAndFileName'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets a list of LoRAs
  Future<Map<String, dynamic>> getLoras() async {
    final response =
        await _httpClient.get(Uri.parse('$host/api/experiment/models/loras'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets details for a specific LoRA
  Future<Map<String, dynamic>> getLoraDetails(String pathAndFileName) async {
    final response = await _httpClient.get(
        Uri.parse('$host/api/view_metadata/loras?filename=$pathAndFileName'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets a list of VAEs
  Future<Map<String, dynamic>> getVaes() async {
    final response =
        await _httpClient.get(Uri.parse('$host/api/experiment/models/vae'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets details for a specific VAE
  Future<Map<String, dynamic>> getVaeDetails(String pathAndFileName) async {
    final response = await _httpClient.get(
        Uri.parse('$host/api/view_metadata/vae?filename=$pathAndFileName'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets a list of upscale models
  Future<Map<String, dynamic>> getUpscaleModels() async {
    final response = await _httpClient
        .get(Uri.parse('$host/api/experiment/models/upscale_models'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets details for a specific upscale model
  Future<Map<String, dynamic>> getUpscaleModelDetails(
      String pathAndFileName) async {
    final response = await _httpClient.get(Uri.parse(
        '$host/api/view_metadata/upscale_models?filename=$pathAndFileName'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets a list of embeddings
  Future<Map<String, dynamic>> getEmbeddings() async {
    final response = await _httpClient
        .get(Uri.parse('$host/api/experiment/models/embeddings'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets details for a specific embedding
  Future<Map<String, dynamic>> getEmbeddingDetails(
      String pathAndFileName) async {
    final response = await _httpClient.get(Uri.parse(
        '$host/api/view_metadata/embeddings?filename=$pathAndFileName'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Gets information about all available objects (nodes)
  Future<Map<String, dynamic>> getObjectInfo() async {
    final response = await _httpClient.get(Uri.parse('$host/api/object_info'));
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Submits a prompt (workflow) to generate an image
  Future<Map<String, dynamic>> submitPrompt(Map<String, dynamic> prompt) async {
    final response = await _httpClient.post(
      Uri.parse('$host/api/prompt'),
      headers: {'Content-Type': 'application/json'},
      body: jsonEncode(prompt),
    );
    _validateResponse(response);
    return jsonDecode(response.body);
  }

  /// Validates HTTP response and throws an exception if needed
  void _validateResponse(http.Response response) {
    if (response.statusCode < 200 || response.statusCode >= 300) {
      throw ComfyUiApiException(
          statusCode: response.statusCode,
          message: 'API request failed: ${response.body}');
    }
  }
}

/// Exception thrown when the ComfyUI API returns an error
class ComfyUiApiException implements Exception {
  final int statusCode;
  final String message;

  ComfyUiApiException({required this.statusCode, required this.message});

  @override
  String toString() => 'ComfyUiApiException: $statusCode - $message';
}