@sogni-ai/sogni-client 4.0.0-alpha.2 → 4.0.0-alpha.20

package/src/ApiClient/WebSocketClient/BrowserWebSocketClient/index.ts

@@ -0,0 +1,241 @@
+import { IWebSocketClient, SupernetType } from '../types';
+import { AuthManager, TokenAuthManager } from '../../../lib/AuthManager';
+import { Logger } from '../../../lib/DefaultLogger';
+import WebSocketClient from '../index';
+import RestClient from '../../../lib/RestClient';
+import { SocketEventMap } from '../events';
+import { MessageType, SocketMessageMap } from '../messages';
+import ChannelCoordinator from './ChannelCoordinator';
+
+interface SocketSend<T extends MessageType = MessageType> {
+  type: 'socket-send';
+  payload: { type: T; data: SocketMessageMap[T] };
+}
+
+interface SocketConnect {
+  type: 'connect';
+}
+
+interface SocketDisconnect {
+  type: 'disconnect';
+}
+
+interface SwitchNetwork {
+  type: 'switchNetwork';
+  payload: SupernetType;
+}
+
+type Message = SocketConnect | SocketDisconnect | SocketSend | SwitchNetwork;
+
+interface EventNotification<T extends keyof SocketEventMap = keyof SocketEventMap> {
+  type: 'socket-event';
+  payload: { type: T; data: SocketEventMap[T] };
+}
+
+interface AuthStateChanged {
+  type: 'auth-state-changed';
+  payload: boolean;
+}
+
+type Notification = EventNotification | AuthStateChanged;
+
+type EventInterceptor<T extends keyof SocketEventMap = keyof SocketEventMap> = (
+  eventType: T,
+  payload: SocketEventMap[T]
+) => void;
+
+class WrappedClient extends WebSocketClient {
+  private interceptor: EventInterceptor | undefined = undefined;
+  intercept(interceptor: EventInterceptor) {
+    this.interceptor = interceptor;
+  }
+  protected emit<T extends keyof SocketEventMap>(event: T, data: SocketEventMap[T]) {
+    super.emit(event, data);
+    if (this.interceptor) {
+      this.interceptor(event, data);
+    }
+  }
+}
+
+class BrowserWebSocketClient extends RestClient<SocketEventMap> implements IWebSocketClient {
+  appId: string;
+  baseUrl: string;
+  private socketClient: WrappedClient;
+  private coordinator: ChannelCoordinator<Message, Notification>;
+  private _isConnected = false;
+  private _supernetType: SupernetType;
+
+  constructor(
+    baseUrl: string,
+    auth: AuthManager,
+    appId: string,
+    supernetType: SupernetType,
+    logger: Logger
+  ) {
+    const socketClient = new WrappedClient(baseUrl, auth, appId, supernetType, logger);
+    super(socketClient.baseUrl, auth, logger);
+    this.socketClient = socketClient;
+    this.appId = appId;
+    this.baseUrl = socketClient.baseUrl;
+    this._supernetType = supernetType;
+    this.coordinator = new ChannelCoordinator({
+      callbacks: {
+        onRoleChange: this.handleRoleChange.bind(this),
+        onMessage: this.handleMessage.bind(this),
+        onNotification: this.handleNotification.bind(this)
+      },
+      logger
+    });
+    this.auth.on('updated', this.handleAuthUpdated.bind(this));
+    this.socketClient.intercept(this.handleSocketEvent.bind(this));
+    //@ts-expect-error window is defined in browser
+    window.DISCONNECT = () => {
+      this.disconnect();
+    };
+  }
+
+  get isConnected() {
+    return this.coordinator.isPrimary ? this.socketClient.isConnected : this._isConnected;
+  }
+
+  get supernetType() {
+    return this.coordinator.isPrimary ? this.socketClient.supernetType : this._supernetType;
+  }
+
+  async connect(): Promise<void> {
+    await this.coordinator.isReady();
+    if (this.coordinator.isPrimary) {
+      await this.socketClient.connect();
+    } else {
+      return this.coordinator.sendMessage({
+        type: 'connect'
+      });
+    }
+  }
+
+  async disconnect() {
+    await this.coordinator.isReady();
+    if (this.coordinator.isPrimary) {
+      this.socketClient.disconnect();
+    } else {
+      this.coordinator.sendMessage({
+        type: 'disconnect'
+      });
+    }
+  }
+
+  async switchNetwork(supernetType: SupernetType): Promise<SupernetType> {
+    await this.coordinator.isReady();
+    if (this.coordinator.isPrimary) {
+      return this.socketClient.switchNetwork(supernetType);
+    }
+    await this.coordinator.sendMessage({
+      type: 'switchNetwork',
+      payload: supernetType
+    });
+    this._supernetType = supernetType;
+    return supernetType;
+  }
+
+  async send<T extends MessageType>(messageType: T, data: SocketMessageMap[T]): Promise<void> {
+    await this.coordinator.isReady();
+    if (this.coordinator.isPrimary) {
+      if (!this.socketClient.isConnected) {
+        await this.socketClient.connect();
+      }
+      return this.socketClient.send(messageType, data);
+    }
+    return this.coordinator.sendMessage({
+      type: 'socket-send',
+      payload: { type: messageType, data }
+    });
+  }
+
+  private async handleMessage(message: Message) {
+    this._logger.debug('Received control message', message);
+    switch (message.type) {
+      case 'socket-send': {
+        if (!this.socketClient.isConnected) {
+          await this.socketClient.connect();
+        }
+        return this.socketClient.send(message.payload.type, message.payload.data);
+      }
+      case 'connect': {
+        if (!this.socketClient.isConnected) {
+          await this.socketClient.connect();
+        }
+        return;
+      }
+      case 'disconnect': {
+        if (this.socketClient.isConnected) {
+          this.socketClient.disconnect();
+        }
+        return;
+      }
+      case 'switchNetwork': {
+        await this.switchNetwork(message.payload);
+        return;
+      }
+      default: {
+        this._logger.error('Received unknown message type:', message);
+      }
+    }
+  }
+
+  private async handleNotification(notification: Notification) {
+    this._logger.debug('Received notification', notification.type, notification.payload);
+    switch (notification.type) {
+      case 'socket-event': {
+        this.emit(notification.payload.type, notification.payload.data);
+        return;
+      }
+      case 'auth-state-changed': {
+        this.handleAuthChanged(notification.payload);
+        return;
+      }
+      default: {
+        this._logger.error('Received unknown notification type:', notification);
+      }
+    }
+  }
+
+  private handleAuthChanged(isAuthenticated: boolean) {
+    if (this.auth instanceof TokenAuthManager) {
+      throw new Error('TokenAuthManager is not supported in multi client mode');
+    }
+    if (this.auth.isAuthenticated !== isAuthenticated) {
+      if (isAuthenticated) {
+        this.auth.authenticate();
+      } else {
+        this.auth.clear();
+      }
+    }
+  }
+
+  private handleSocketEvent(eventType: keyof SocketEventMap, payload: any) {
+    if (this.coordinator.isPrimary) {
+      this.coordinator.notify({
+        type: 'socket-event',
+        payload: { type: eventType, data: payload }
+      });
+      this.emit(eventType, payload);
+    }
+  }
+
+  private handleAuthUpdated(isAuthenticated: boolean) {
+    this.coordinator.notify({
+      type: 'auth-state-changed',
+      payload: isAuthenticated
+    });
+  }
+
+  private handleRoleChange(isPrimary: boolean) {
+    if (isPrimary && !this.socketClient.isConnected && this.isConnected) {
+      this.socketClient.connect();
+    } else if (!isPrimary && this.socketClient.isConnected) {
+      this.socketClient.disconnect();
+    }
+  }
+}
+
+export default BrowserWebSocketClient;

package/src/ApiClient/WebSocketClient/events.ts

@@ -148,3 +148,5 @@ export type SocketEventMap = {
 
   artistCancelConfirmation: ArtistCancelConfirmation;
 };
+
+export type SocketEventName = keyof SocketEventMap;

package/src/ApiClient/WebSocketClient/index.ts

@@ -1,7 +1,7 @@
 import { MessageType, SocketMessageMap } from './messages';
 import { SocketEventMap } from './events';
 import RestClient from '../../lib/RestClient';
-import { SupernetType } from './types';
+import { IWebSocketClient, SupernetType } from './types';
 import WebSocket, { CloseEvent, ErrorEvent, MessageEvent } from 'isomorphic-ws';
 import { base64Decode, base64Encode } from '../../lib/base64';
 import isNodejs from '../../lib/isNodejs';
@@ -13,7 +13,7 @@ const PROTOCOL_VERSION = '3.0.0';
 
 const PING_INTERVAL = 15000;
 
-class WebSocketClient extends RestClient<SocketEventMap> {
+class WebSocketClient extends RestClient<SocketEventMap> implements IWebSocketClient {
   appId: string;
   baseUrl: string;
   private socket: WebSocket | null = null;
@@ -86,7 +86,7 @@ class WebSocketClient extends RestClient<SocketEventMap> {
     socket.onmessage = null;
     socket.onopen = null;
     this.stopPing();
-    socket.close();
+    socket.close(1000, 'Client disconnected');
   }
 
   private startPing(socket: WebSocket) {
@@ -148,7 +148,7 @@ class WebSocketClient extends RestClient<SocketEventMap> {
   }
 
   private handleClose(e: CloseEvent) {
-    if (e.target === this.socket) {
+    if (e.target === this.socket || !this.socket) {
       this._logger.info('WebSocket disconnected, cleanup', e);
       this.disconnect();
       this.emit('disconnected', {

package/src/ApiClient/WebSocketClient/types.ts

@@ -1 +1,17 @@
+import { MessageType, SocketMessageMap } from './messages';
+import RestClient from '../../lib/RestClient';
+import { SocketEventMap } from './events';
+
 export type SupernetType = 'relaxed' | 'fast';
+
+export interface IWebSocketClient extends RestClient<SocketEventMap> {
+  appId: string;
+  baseUrl: string;
+  isConnected: boolean;
+  supernetType: SupernetType;
+
+  connect(): Promise<void>;
+  disconnect(): void;
+  send<T extends MessageType>(messageType: T, data: SocketMessageMap[T]): Promise<void>;
+  switchNetwork(supernetType: SupernetType): Promise<SupernetType>;
+}

package/src/ApiClient/index.ts

@@ -3,12 +3,14 @@ import WebSocketClient from './WebSocketClient';
 import TypedEventEmitter from '../lib/TypedEventEmitter';
 import { ApiClientEvents } from './events';
 import { ServerConnectData, ServerDisconnectData } from './WebSocketClient/events';
-import { isNotRecoverable } from './WebSocketClient/ErrorCode';
+import { ErrorCode, isNotRecoverable } from './WebSocketClient/ErrorCode';
 import { JSONValue } from '../types/json';
-import { SupernetType } from './WebSocketClient/types';
+import { IWebSocketClient, SupernetType } from './WebSocketClient/types';
 import { Logger } from '../lib/DefaultLogger';
 import CookieAuthManager from '../lib/AuthManager/CookieAuthManager';
 import { AuthManager, TokenAuthManager } from '../lib/AuthManager';
+import isNodejs from '../lib/isNodejs';
+import BrowserWebSocketClient from './WebSocketClient/BrowserWebSocketClient';
 
 const WS_RECONNECT_ATTEMPTS = 5;
 
@@ -42,13 +44,14 @@ export interface ApiClientOptions {
   logger: Logger;
   authType: 'token' | 'cookies';
   disableSocket?: boolean;
+  multiInstance?: boolean;
 }
 
 class ApiClient extends TypedEventEmitter<ApiClientEvents> {
   readonly appId: string;
   readonly logger: Logger;
   private _rest: RestClient;
-  private _socket: WebSocketClient;
+  private _socket: IWebSocketClient;
   private _auth: AuthManager;
   private _reconnectAttempts = WS_RECONNECT_ATTEMPTS;
   private _disableSocket: boolean = false;
@@ -60,7 +63,8 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
     networkType,
     authType,
     logger,
-    disableSocket = false
+    disableSocket = false,
+    multiInstance = false
   }: ApiClientOptions) {
     super();
     this.appId = appId;
@@ -68,7 +72,13 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
     this._auth =
       authType === 'token' ? new TokenAuthManager(baseUrl, logger) : new CookieAuthManager(logger);
     this._rest = new RestClient(baseUrl, this._auth, logger);
-    this._socket = new WebSocketClient(socketUrl, this._auth, appId, networkType, logger);
+    const supportMultiInstance = !isNodejs && this._auth instanceof CookieAuthManager;
+    if (supportMultiInstance && multiInstance) {
+      // Use coordinated WebSocket client to share single connection between tabs
+      this._socket = new BrowserWebSocketClient(socketUrl, this._auth, appId, networkType, logger);
+    } else {
+      this._socket = new WebSocketClient(socketUrl, this._auth, appId, networkType, logger);
+    }
     this._disableSocket = disableSocket;
     this._auth.on('updated', this.handleAuthUpdated.bind(this));
     this._socket.on('connected', this.handleSocketConnect.bind(this));
@@ -83,7 +93,7 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
     return this._auth;
   }
 
-  get socket(): WebSocketClient {
+  get socket(): IWebSocketClient {
     return this._socket;
   }
 
@@ -101,14 +111,26 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
   }
 
   handleSocketDisconnect(data: ServerDisconnectData) {
+    // If user is not authenticated, we don't need to reconnect
+    if (!this.auth.isAuthenticated || data.code === 1000) {
+      this.emit('disconnected', data);
+      return;
+    }
     if (!data.code || isNotRecoverable(data.code)) {
+      // If this is browser, another tab is probably claiming the connection, so we don't need to reconnect
+      if (
+        this._socket instanceof BrowserWebSocketClient &&
+        data.code === ErrorCode.SWITCH_CONNECTION
+      ) {
+        this.logger.debug('Switching network connection, not reconnecting');
+        return;
+      }
       this.auth.clear();
       this.emit('disconnected', data);
       this.logger.error('Not recoverable socket error', data);
       return;
     }
     if (this._reconnectAttempts <= 0) {
-      this.auth.clear();
       this.emit('disconnected', data);
       this._reconnectAttempts = WS_RECONNECT_ATTEMPTS;
       return;

package/src/Projects/createJobRequestMessage.ts

@@ -1,6 +1,12 @@
 import { ProjectParams } from './types';
 import { ControlNetParams, ControlNetParamsRaw } from './types/ControlNetParams';
-import { validateNumber, validateCustomImageSize } from '../lib/validation';
+import {
+  validateNumber,
+  validateCustomImageSize,
+  validateSampler,
+  validateScheduler
+} from '../lib/validation';
+
 // Mac worker can't process the data if some of the fields are missing, so we need to provide a default template
 function getTemplate() {
   return {
@@ -25,8 +31,8 @@ function getTemplate() {
         guidanceScaleIsEnabled: true,
         siImageBackgroundColor: 'black',
         cnDragOffset: [0, 0],
-        scheduler: 'DPM Solver Multistep (DPM-Solver++)',
-        timeStepSpacing: 'Linear',
+        scheduler: null,
+        timeStepSpacing: null,
         steps: 20,
         cnRotation: 0,
         guidanceScale: 7.5,
@@ -119,8 +125,8 @@ function createJobRequestMessage(id: string, params: ProjectParams) {
     keyFrames: [
       {
         ...template.keyFrames[0],
-        scheduler: params.scheduler || null,
-        timeStepSpacing: params.timeStepSpacing || null,
+        scheduler: validateSampler(params.sampler),
+        timeStepSpacing: validateScheduler(params.scheduler),
         steps: params.steps,
         guidanceScale: params.guidance,
         modelID: params.modelId,

package/src/Projects/index.ts

@@ -4,6 +4,7 @@ import {
   EnhancementStrength,
   EstimateRequest,
   ImageUrlParams,
+  CostEstimation,
   ProjectParams,
   SizePreset,
   SupportedModel
@@ -28,6 +29,7 @@ import Cache from '../lib/Cache';
 import { enhancementDefaults } from './Job';
 import { getEnhacementStrength } from './utils';
 import { TokenType } from '../types/token';
+import { validateSampler } from '../lib/validation';
 
 const sizePresetCache = new Cache<SizePreset[]>(10 * 60 * 1000);
 const GARBAGE_COLLECT_TIMEOUT = 30000;
@@ -515,9 +517,9 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     height,
     sizePreset,
     guidance,
-    scheduler,
+    sampler,
     contextImages
-  }: EstimateRequest) {
+  }: EstimateRequest): Promise<CostEstimation> {
     let apiVersion = 2;
     const pathParams = [
       tokenType || 'sogni',
@@ -541,10 +543,10 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     } else {
       pathParams.push(0, 0);
     }
-    if (scheduler) {
+    if (sampler) {
       apiVersion = 3;
       pathParams.push(guidance || 0);
-      pathParams.push(scheduler || '');
+      pathParams.push(validateSampler(sampler)!);
       pathParams.push(contextImages || 0);
     }
     const r = await this.client.socket.get<EstimationResponse>(
@@ -552,7 +554,9 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     );
     return {
       token: r.quote.project.costInToken,
-      usd: r.quote.project.costInUSD
+      usd: r.quote.project.costInUSD,
+      spark: r.quote.project.costInSpark,
+      sogni: r.quote.project.costInSogni
     };
   }
 

package/src/Projects/types/EstimationResponse.ts

@@ -14,6 +14,8 @@ export interface Job {
   costInRenderSec: string;
   costInUSD: string;
   costInToken: string;
+  costInSpark: string;
+  costInSogni: string;
   calculatedStepCount?: number;
 }
 

package/src/Projects/types/SamplerParams.ts

@@ -0,0 +1,24 @@
+export const SupportedSamplers = {
+  dfs_sd3: 'Discrete Flow Scheduler (SD3)',
+  dpm_pp: 'DPM Solver Multistep (DPM-Solver++)',
+  dpm_pp_sde: 'DPM++ SDE',
+  dpm_pp_2m: 'DPM++ 2M',
+  //dpm_pp_2m_sde: 'DPM++ 2M SDE',
+  euler: 'Euler',
+  euler_a: 'Euler a',
+  //heun: 'Heun',
+  lcm: 'LCM (Latent Consistency Model)',
+  pndm_plms: 'PNDM (Pseudo-linear multi-step)'
+  //uni_pc: 'UniPC'
+};
+
+export function isSampler(sampler: string): sampler is Sampler {
+  return sampler in SupportedSamplers;
+}
+
+export function isRawSampler(sampler: string): boolean {
+  const samplers = Object.values(SupportedSamplers);
+  return samplers.includes(sampler);
+}
+
+export type Sampler = keyof typeof SupportedSamplers;

package/src/Projects/types/SchedulerParams.ts

@@ -0,0 +1,22 @@
+export const SupportedSchedulers = {
+  beta: 'Beta',
+  ddim: 'DDIM',
+  karras: 'Karras',
+  kl_optimal: 'KL Optimal',
+  leading: 'Leading',
+  linear: 'Linear',
+  normal: 'Normal',
+  sgm_uniform: 'SGM Uniform',
+  simple: 'Simple'
+};
+
+export function isScheduler(scheduler: string): scheduler is Scheduler {
+  return scheduler in SupportedSchedulers;
+}
+
+export function isRawScheduler(scheduler: string): boolean {
+  const schedulers = Object.values(SupportedSchedulers);
+  return schedulers.includes(scheduler);
+}
+
+export type Scheduler = keyof typeof SupportedSchedulers;

package/src/Projects/types/index.ts

@@ -1,6 +1,8 @@
 import { SupernetType } from '../../ApiClient/WebSocketClient/types';
 import { ControlNetParams } from './ControlNetParams';
 import { TokenType } from '../../types/token';
+import { Sampler, SupportedSamplers } from './SamplerParams';
+import { Scheduler, SupportedSchedulers } from './SchedulerParams';
 
 export interface SupportedModel {
   id: string;
@@ -23,15 +25,9 @@ export interface SizePreset {
   aspect: string;
 }
 
-export type Scheduler =
-  | 'DPM Solver Multistep (DPM-Solver++)'
-  | 'PNDM (Pseudo-linear multi-step)'
-  | 'LCM (Latent Consistency Model)'
-  | 'Discrete Flow Scheduler (SD3)'
-  | 'Euler' // Used by Flux
-  | 'Euler a'; // Used by Flux
+export type { Sampler, Scheduler };
 
-export type TimeStepSpacing = 'Leading' | 'Linear' | 'Karras' | 'Simple' | 'SGM Uniform';
+export { SupportedSamplers, SupportedSchedulers };
 
 export type OutputFormat = 'png' | 'jpg';
 
@@ -108,11 +104,11 @@ export interface ProjectParams {
   /**
    * Scheduler to use
    */
-  scheduler?: Scheduler;
+  sampler?: Sampler;
   /**
    * Time step spacing method
    */
-  timeStepSpacing?: TimeStepSpacing;
+  scheduler?: Scheduler;
   /**
    * Size preset ID to use. You can query available size presets
    * from `client.projects.sizePresets(network, modelId)`
@@ -202,9 +198,9 @@ export interface EstimateRequest {
    */
   guidance?: number;
   /**
-   * Scheduler
+   * Sampler
    */
-  scheduler?: Scheduler;
+  sampler?: Sampler;
   /**
    * Number of context images to use (for Flux Kontext).
    * Note that this parameter is ignored if `scheduler` is not provided
@@ -212,4 +208,18 @@ export interface EstimateRequest {
   contextImages?: number;
 }
 
+/**
+ * Represents estimation of project cost in different currency formats
+ */
+export interface CostEstimation {
+  /** Cost in selected token type */
+  token: string;
+  /** Cost in USD */
+  usd: string;
+  /** Cost in Spark Points */
+  spark: string;
+  /** Cost in Sogni tokens */
+  sogni: string;
+}
+
 export type EnhancementStrength = 'light' | 'medium' | 'heavy';

package/src/index.ts

@@ -16,8 +16,10 @@ import {
   AvailableModel,
   OutputFormat,
   ProjectParams,
+  Sampler,
   Scheduler,
-  TimeStepSpacing
+  SupportedSamplers,
+  SupportedSchedulers
 } from './Projects/types';
 // Stats API
 import StatsApi from './Stats';
@@ -36,13 +38,13 @@ export type {
   OutputFormat,
   ProjectParams,
   ProjectStatus,
-  Scheduler,
+  Sampler,
   SupernetType,
-  TimeStepSpacing,
+  Scheduler,
   TokenType
 };
 
-export { ApiError, CurrentAccount, Job, Project };
+export { ApiError, CurrentAccount, Job, Project, SupportedSamplers, SupportedSchedulers };
 
 export interface SogniClientConfig {
   /**
@@ -61,7 +63,7 @@ export interface SogniClientConfig {
   socketEndpoint?: string;
   /**
    * Disable WebSocket connection. Useful for testing or when WebSocket is not needed.
-   * Note that many may not work without WebSocket connection.
+   * Note that many APIs may not work without WebSocket connection.
    * @experimental
    * @internal
    */
@@ -93,6 +95,14 @@ export interface SogniClientConfig {
    * @experimental
    */
   authType?: 'token' | 'cookies';
+  /**
+   * Browser only. If true, the client will use a single WebSocket connection shared across multiple
+   * tabs. This is useful for browser apps that need to process multiple projects at the same time.
+   * Only works in browser environment and with cookie authentication.
+   * @default false
+   * @experimental
+   */
+  multiInstance?: boolean;
 }
 
 export class SogniClient {
@@ -139,9 +149,13 @@ export class SogniClient {
       throw Error('This method should only be called when using cookie auth');
     }
     try {
-      await this.apiClient.rest.get<ApiResponse<MeData>>('/v1/account/me');
+      const res = await this.apiClient.rest.get<ApiResponse<MeData>>('/v1/account/me');
       await auth.authenticate();
-      await this.account.me();
+      this.currentAccount._update({
+        username: res.data.username,
+        email: res.data.currentEmail,
+        walletAddress: res.data.walletAddress
+      });
       return true;
     } catch (e) {
       this.apiClient.logger.info('Client is not authenticated');
@@ -167,7 +181,8 @@ export class SogniClient {
       networkType: network,
       logger,
       authType: config.authType || 'token',
-      disableSocket: config.disableSocket
+      disableSocket: config.disableSocket,
+      multiInstance: config.multiInstance
     });
     const eip712 = new EIP712Helper({
       name: isTestnet ? 'Sogni-testnet' : 'Sogni AI',