@sogni-ai/sogni-client 1.0.2 → 1.1.0

package/src/ApiClient/WebSocketClient/index.ts

@@ -1,13 +1,15 @@
 import { MessageType, SocketMessageMap } from './messages';
 import { SocketEventMap } from './events';
-import RestClient, { AuthData } from '../../lib/RestClient';
+import RestClient from '../../lib/RestClient';
 import { SupernetType } from './types';
 import WebSocket, { CloseEvent, ErrorEvent, MessageEvent } from 'isomorphic-ws';
 import { base64Decode, base64Encode } from '../../lib/base64';
 import isNodejs from '../../lib/isNodejs';
-import Cookie from 'js-cookie';
 import { LIB_VERSION } from '../../version';
 import { Logger } from '../../lib/DefaultLogger';
+import AuthManager from '../../lib/AuthManager';
+
+const PROTOCOL_VERSION = '0.4.3';
 
 const PING_INTERVAL = 15000;
 
@@ -18,34 +20,23 @@ class WebSocketClient extends RestClient<SocketEventMap> {
   private _supernetType: SupernetType;
   private _pingInterval: NodeJS.Timeout | null = null;
 
-  constructor(baseUrl: string, appId: string, supernetType: SupernetType, logger: Logger) {
+  constructor(
+    baseUrl: string,
+    auth: AuthManager,
+    appId: string,
+    supernetType: SupernetType,
+    logger: Logger
+  ) {
     const _baseUrl = new URL(baseUrl);
     if (_baseUrl.protocol === 'wss:') {
       _baseUrl.protocol = 'https:';
     }
-    super(_baseUrl.toString(), logger);
+    super(_baseUrl.toString(), auth, logger);
     this.appId = appId;
     this.baseUrl = _baseUrl.toString();
     this._supernetType = supernetType;
   }
 
-  set auth(auth: AuthData | null) {
-    //In browser, set the cookie
-    if (!isNodejs) {
-      if (auth) {
-        Cookie.set('authorization', auth.token, {
-          domain: '.sogni.ai',
-          expires: 1
-        });
-      } else {
-        Cookie.remove('authorization', {
-          domain: '.sogni.ai'
-        });
-      }
-    }
-    this._auth = auth;
-  }
-
   get supernetType(): SupernetType {
     return this._supernetType;
   }
@@ -54,11 +45,11 @@ class WebSocketClient extends RestClient<SocketEventMap> {
     return !!this.socket;
   }
 
-  connect() {
+  async connect() {
     if (this.socket) {
       this.disconnect();
     }
-    const userAgent = `Sogni/${LIB_VERSION} (sogni-client)`;
+    const userAgent = `Sogni/${PROTOCOL_VERSION} (sogni-client) ${LIB_VERSION}`;
     const url = new URL(this.baseUrl);
     url.protocol = 'wss:';
     url.searchParams.set('appId', this.appId);
@@ -71,7 +62,7 @@ class WebSocketClient extends RestClient<SocketEventMap> {
     if (isNodejs) {
       params = {
         headers: {
-          Authorization: this._auth?.token,
+          Authorization: await this.auth.getToken(),
           'User-Agent': userAgent
         }
       };
@@ -114,7 +105,7 @@ class WebSocketClient extends RestClient<SocketEventMap> {
   }
 
   switchNetwork(supernetType: SupernetType): Promise<SupernetType> {
-    return new Promise<SupernetType>(async (resolve, reject) => {
+    return new Promise<SupernetType>(async (resolve) => {
       this.once('changeNetwork', ({ network }) => {
         this._supernetType = network;
         resolve(network);

package/src/ApiClient/WebSocketClient/messages.ts

@@ -1,8 +1,16 @@
 import { JobRequestRaw } from '../../Projects/createJobRequestMessage';
 import { SupernetType } from './types';
 
+export interface JobErrorMessage {
+  jobID: string;
+  error: 'artistCanceled';
+  error_message: 'artistCanceled';
+  isFromWorker: false;
+}
+
 export interface SocketMessageMap {
   jobRequest: JobRequestRaw;
+  jobError: JobErrorMessage;
   changeNetwork: SupernetType;
 }
 

package/src/ApiClient/index.ts

@@ -1,6 +1,5 @@
 import RestClient from '../lib/RestClient';
 import WebSocketClient from './WebSocketClient';
-import { jwtDecode } from 'jwt-decode';
 import TypedEventEmitter from '../lib/TypedEventEmitter';
 import { ApiClientEvents } from './events';
 import { ServerConnectData, ServerDisconnectData } from './WebSocketClient/events';
@@ -8,6 +7,7 @@ import { isNotRecoverable } from './WebSocketClient/ErrorCode';
 import { JSONValue } from '../types/json';
 import { SupernetType } from './WebSocketClient/types';
 import { Logger } from '../lib/DefaultLogger';
+import AuthManager, { Tokens } from '../lib/AuthManager';
 
 const WS_RECONNECT_ATTEMPTS = 5;
 
@@ -33,21 +33,12 @@ export class ApiError extends Error {
   }
 }
 
-/**
- * @inline
- */
-interface AuthData {
-  token: string;
-  walletAddress: string;
-  expiresAt: Date;
-}
-
 class ApiClient extends TypedEventEmitter<ApiClientEvents> {
   readonly appId: string;
   readonly logger: Logger;
   private _rest: RestClient;
   private _socket: WebSocketClient;
-  private _auth: AuthData | null = null;
+  private _auth: AuthManager;
   private _reconnectAttempts = WS_RECONNECT_ATTEMPTS;
 
   constructor(
@@ -60,18 +51,21 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
     super();
     this.appId = appId;
     this.logger = logger;
-    this._rest = new RestClient(baseUrl, logger);
-    this._socket = new WebSocketClient(socketUrl, appId, networkType, logger);
+    this._auth = new AuthManager(baseUrl, logger);
+    this._rest = new RestClient(baseUrl, this._auth, logger);
+    this._socket = new WebSocketClient(socketUrl, this._auth, appId, networkType, logger);
+
+    this._auth.on('refreshFailed', this.handleRefreshFailed.bind(this));
     this._socket.on('connected', this.handleSocketConnect.bind(this));
     this._socket.on('disconnected', this.handleSocketDisconnect.bind(this));
   }
 
   get isAuthenticated(): boolean {
-    return !!this._auth && this._auth.expiresAt > new Date();
+    return this.auth.isAuthenticated;
   }
 
-  get auth(): AuthData | null {
-    return this._auth && this._auth.expiresAt > new Date() ? this._auth : null;
+  get auth(): AuthManager {
+    return this._auth;
   }
 
   get socket(): WebSocketClient {
@@ -82,20 +76,13 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
     return this._rest;
   }
 
-  authenticate(token: string) {
-    const decoded = jwtDecode<{ addr: string; env: string; iat: number; exp: number }>(token);
-    this._auth = {
-      token,
-      walletAddress: decoded.addr,
-      expiresAt: new Date(decoded.exp * 1000)
-    };
-    this.rest.auth = { token };
-    this.socket.auth = { token };
-    this.socket.connect();
+  async authenticate(tokens: Tokens) {
+    await this.auth.setTokens(tokens);
+    await this.socket.connect();
   }
 
   removeAuth() {
-    this._auth = null;
+    this.auth.clear();
     this.socket.disconnect();
   }
 
@@ -112,6 +99,7 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
       return;
     }
     if (this._reconnectAttempts <= 0) {
+      this.removeAuth();
       this.emit('disconnected', data);
       this._reconnectAttempts = WS_RECONNECT_ATTEMPTS;
       return;
@@ -119,6 +107,10 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
     this._reconnectAttempts--;
     setTimeout(() => this.socket.connect(), 1000);
   }
+
+  handleRefreshFailed() {
+    this.removeAuth();
+  }
 }
 
 export default ApiClient;

package/src/Projects/Project.ts

@@ -164,6 +164,13 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
     });
   }
 
+  /**
+   * Cancel the project. This will cancel all jobs in the project.
+   */
+  async cancel() {
+    await this._api.cancel(this.id);
+  }
+
   /**
    * Find a job by id
    * @param id

package/src/Projects/createJobRequestMessage.ts

@@ -65,6 +65,14 @@ function getTemplate() {
   };
 }
 
+function validateSize(value: any): number {
+  const size = Number(value);
+  if (isNaN(size) || size < 256 || size > 2048) {
+    throw new Error('Width and height must be numbers between 256 and 2048');
+  }
+  return size;
+}
+
 function createJobRequestMessage(id: string, params: ProjectParams) {
   const template = getTemplate();
   const jobRequest: Record<string, any> = {
@@ -84,7 +92,8 @@ function createJobRequestMessage(id: string, params: ProjectParams) {
         strengthIsEnabled: !!params.startingImage,
         strength: !!params.startingImage
           ? 1 - (Number(params.startingImageStrength) || 0.5)
-          : undefined
+          : undefined,
+        sizePreset: params.sizePreset
       }
     ],
     previews: params.numberOfPreviews || 0,
@@ -95,6 +104,10 @@ function createJobRequestMessage(id: string, params: ProjectParams) {
   if (params.network) {
     jobRequest.network = params.network;
   }
+  if (params.sizePreset === 'custom') {
+    jobRequest.keyFrames[0].width = validateSize(params.width);
+    jobRequest.keyFrames[0].height = validateSize(params.height);
+  }
   return jobRequest;
 }
 

package/src/Projects/index.ts

@@ -1,6 +1,12 @@
 import ApiGroup, { ApiConfig } from '../ApiGroup';
-import models from './models.json';
-import { AvailableModel, EstimateRequest, ImageUrlParams, ProjectParams } from './types';
+import {
+  AvailableModel,
+  EstimateRequest,
+  ImageUrlParams,
+  ProjectParams,
+  SizePreset,
+  SupportedModel
+} from './types';
 import {
   JobErrorData,
   JobProgressData,
@@ -16,8 +22,12 @@ import { JobEvent, ProjectApiEvents, ProjectEvent } from './types/events';
 import getUUID from '../lib/getUUID';
 import { RawProject } from './types/RawProject';
 import ErrorData from '../types/ErrorData';
+import { SupernetType } from '../ApiClient/WebSocketClient/types';
+import Cache from '../lib/Cache';
 
+const sizePresetCache = new Cache<SizePreset[]>(10 * 60 * 1000);
 const GARBAGE_COLLECT_TIMEOUT = 10000;
+const MODELS_REFRESH_INTERVAL = 1000 * 60 * 60 * 24; // 24 hours
 
 function mapErrorCodes(code: string): number {
   switch (code) {
@@ -39,6 +49,10 @@ function mapErrorCodes(code: string): number {
 class ProjectsApi extends ApiGroup<ProjectApiEvents> {
   private _availableModels: AvailableModel[] = [];
   private projects: Project[] = [];
+  private _supportedModels: { data: SupportedModel[] | null; updatedAt: Date } = {
+    data: null,
+    updatedAt: new Date(0)
+  };
 
   get availableModels() {
     return this._availableModels;
@@ -65,14 +79,20 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     this.emit('availableModels', this._availableModels);
   }
 
-  private handleSwarmModels(data: SocketEventMap['swarmModels']) {
-    const modelIndex = models.reduce((acc: Record<string, any>, model) => {
-      acc[model.modelId] = model;
+  private async handleSwarmModels(data: SocketEventMap['swarmModels']) {
+    let models: SupportedModel[] = [];
+    try {
+      models = await this.getSupportedModels();
+    } catch (e) {
+      this.client.logger.error(e);
+    }
+    const modelIndex = models.reduce((acc: Record<string, SupportedModel>, model) => {
+      acc[model.id] = model;
       return acc;
     }, {});
     this._availableModels = Object.entries(data).map(([id, workerCount]) => ({
       id,
-      name: modelIndex[id]?.modelShortName || id.replace(/-/g, ' '),
+      name: modelIndex[id]?.name || id.replace(/-/g, ' '),
       workerCount
     }));
     this.emit('availableModels', this._availableModels);
@@ -339,6 +359,38 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     return data.project;
   }
 
+  /**
+   * Cancel project by id. This will cancel all jobs in the project and mark project as canceled.
+   * Client may still receive job events for the canceled jobs as it takes some time, but they will
+   * be ignored
+   * @param projectId
+   **/
+  async cancel(projectId: string) {
+    await this.client.socket.send('jobError', {
+      jobID: projectId,
+      error: 'artistCanceled',
+      error_message: 'artistCanceled',
+      isFromWorker: false
+    });
+    const project = this.projects.find((p) => p.id === projectId);
+    if (!project) {
+      return;
+    }
+    // Remove project from the list to stop tracking it
+    this.projects = this.projects.filter((p) => p.id !== projectId);
+
+    // Cancel all jobs in the project
+    project.jobs.forEach((job) => {
+      if (!job.finished) {
+        job._update({ status: 'canceled' });
+      }
+    });
+    // If project is still in processing, mark it as canceled
+    if (!project.finished) {
+      project._update({ status: 'canceled' });
+    }
+  }
+
   private async uploadGuideImage(projectId: string, file: File | Buffer | Blob) {
     const imageId = getUUID();
     const presignedUrl = await this.uploadUrl({
@@ -370,10 +422,32 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     stepCount,
     previewCount,
     cnEnabled,
-    startingImageStrength
+    startingImageStrength,
+    width,
+    height,
+    sizePreset
   }: EstimateRequest) {
+    const pathParams = [
+      network,
+      model,
+      imageCount,
+      stepCount,
+      previewCount,
+      cnEnabled ? 1 : 0,
+      startingImageStrength ? 1 - startingImageStrength : 0
+    ];
+    if (sizePreset) {
+      const presets = await this.getSizePresets(network, model);
+      const preset = presets.find((p) => p.id === sizePreset);
+      if (!preset) {
+        throw new Error('Invalid size preset');
+      }
+      pathParams.push(preset.width, preset.height);
+    } else if (width && height) {
+      pathParams.push(width, height);
+    }
     const r = await this.client.socket.get<EstimationResponse>(
-      `/api/v1/job/estimate/${network}/${model}/${imageCount}/${stepCount}/${previewCount}/${cnEnabled ? 1 : 0}/${startingImageStrength ? 1 - startingImageStrength : 0}`
+      `/api/v1/job/estimate/${pathParams.join('/')}`
     );
     return {
       token: r.quote.project.costInToken,
@@ -406,6 +480,74 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     );
     return r.data.downloadUrl;
   }
+
+  async getSupportedModels(forceRefresh = false) {
+    if (
+      this._supportedModels.data &&
+      !forceRefresh &&
+      Date.now() - this._supportedModels.updatedAt.getTime() < MODELS_REFRESH_INTERVAL
+    ) {
+      return this._supportedModels.data;
+    }
+    const models = await this.client.socket.get<SupportedModel[]>(`/api/v1/models/list`);
+    this._supportedModels = { data: models, updatedAt: new Date() };
+    return models;
+  }
+
+  /**
+   * Get supported size presets for the model and network. Size presets are cached for 10 minutes.
+   *
+   * @example
+   * ```ts
+   * const presets = await client.projects.getSizePresets('fast', 'flux1-schnell-fp8');
+   * console.log(presets);
+   * ```
+   *
+   * @param network - 'fast' or 'relaxed'
+   * @param modelId - model id (e.g. 'flux1-schnell-fp8')
+   * @param forceRefresh - force refresh cache
+   * @returns {Promise<{
+   *   label: string;
+   *   id: string;
+   *   width: number;
+   *   height: number;
+   *   ratio: string;
+   *   aspect: string;
+   * }[]>}
+   */
+  async getSizePresets(network: SupernetType, modelId: string, forceRefresh = false) {
+    const key = `${network}-${modelId}`;
+    const cached = sizePresetCache.read(key);
+    if (cached && !forceRefresh) {
+      return cached;
+    }
+    const data = await this.client.socket.get<SizePreset[]>(
+      `/api/v1/size-presets/network/${network}/model/${modelId}`
+    );
+    sizePresetCache.write(key, data);
+    return data;
+  }
+
+  /**
+   * Get available models and their worker counts. Normally, you would get list once you connect
+   * to the server, but you can also call this method to get the list of available models manually.
+   * @param network
+   */
+  async getAvailableModels(network: SupernetType): Promise<AvailableModel[]> {
+    const workersByModelSid = await this.client.socket.get<Record<string, number>>(
+      `/api/v1/status/network/${network}/models`
+    );
+    const supportedModels = await this.getSupportedModels();
+    return Object.entries(workersByModelSid).map(([sid, workerCount]) => {
+      const SID = Number(sid);
+      const model = supportedModels.find((m) => m.SID === SID);
+      return {
+        id: model?.id || sid,
+        name: model?.name || sid.replace(/-/g, ' '),
+        workerCount
+      };
+    });
+  }
 }
 
 export default ProjectsApi;

package/src/Projects/types/index.ts

@@ -1,11 +1,26 @@
 import { SupernetType } from '../../ApiClient/WebSocketClient/types';
 
+export interface SupportedModel {
+  id: string;
+  name: string;
+  SID: number;
+}
+
 export interface AvailableModel {
   id: string;
   name: string;
   workerCount: number;
 }
 
+export interface SizePreset {
+  label: string;
+  id: string;
+  width: number;
+  height: number;
+  ratio: string;
+  aspect: string;
+}
+
 export interface AiModel {
   isSD3: boolean;
   modelShortName: string;
@@ -108,6 +123,19 @@ export interface ProjectParams {
    * Time step spacing method
    */
   timeStepSpacing?: TimeStepSpacing;
+  /**
+   * Size preset ID to use. You can query available size presets
+   * from `client.projects.sizePresets(network, modelId)`
+   */
+  sizePreset?: 'custom' | string;
+  /**
+   * Output image width. Only used if `sizePreset` is "custom"
+   */
+  width?: number;
+  /**
+   * Output image height. Only used if `sizePreset` is "custom"
+   */
+  height?: number;
 }
 
 export type ImageUrlParams = {
@@ -147,4 +175,18 @@ export interface EstimateRequest {
    * How strong effect of starting image should be. From 0 to 1, default 0.5
    */
   startingImageStrength?: number;
+  /**
+   * Size preset ID
+   */
+  sizePreset?: string;
+  /**
+   * Size preset image width, if not using size preset
+   * @internal
+   */
+  width?: number;
+  /**
+   * Size preset image height, if not using size preset
+   * @internal
+   */
+  height?: number;
 }