@sogni-ai/sogni-client 4.0.0-alpha.3 → 4.0.0-alpha.30

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,9 +148,16 @@ class WebSocketClient extends RestClient<SocketEventMap> {
   }
 
   private handleClose(e: CloseEvent) {
-    if (e.target === this.socket) {
+    const socket = e.target;
+    socket.onerror = null;
+    socket.onmessage = null;
+    socket.onopen = null;
+    if (socket === this.socket || !this.socket) {
       this._logger.info('WebSocket disconnected, cleanup', e);
-      this.disconnect();
+      if (socket === this.socket) {
+        this.stopPing();
+        this.socket = null;
+      }
       this.emit('disconnected', {
         code: e.code,
         reason: e.reason
@@ -193,6 +200,9 @@ class WebSocketClient extends RestClient<SocketEventMap> {
   }
 
   async send<T extends MessageType>(messageType: T, data: SocketMessageMap[T]) {
+    if (!this.isConnected) {
+      await this.connect();
+    }
     await this.waitForConnection();
     this._logger.debug('WebSocket send:', messageType, data);
     this.socket!.send(

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;
@@ -122,7 +144,7 @@ class ApiClient extends TypedEventEmitter<ApiClientEvents> {
       if (this.socket.isConnected) {
         this.socket.disconnect();
       }
-    } else if (!this._disableSocket) {
+    } else if (!this._disableSocket && !this.socket.isConnected) {
       this.socket.connect();
     }
   }

package/src/Projects/Job.ts

@@ -19,7 +19,7 @@ export const enhancementDefaults = {
   startingImageStrength: 0.5,
   steps: 5,
   guidance: 1,
-  numberOfImages: 1,
+  numberOfMedia: 1,
   numberOfPreviews: 0
 };
 
@@ -61,6 +61,11 @@ export interface JobData {
   positivePrompt?: string;
   negativePrompt?: string;
   jobIndex?: number;
+  /**
+   * Estimated time remaining in seconds (for long-running jobs like video generation).
+   * Updated by ComfyUI workers during inference.
+   */
+  etaSeconds?: number;
 }
 
 export interface JobEventMap extends EntityEvents {
@@ -178,10 +183,21 @@ class Job extends DataEntity<JobData, JobEventMap> {
     return this.data.error;
   }
 
-  get hasResultImage() {
+  /**
+   * Whether this job has a result media file available for download.
+   * Returns true if completed and not NSFW filtered.
+   */
+  get hasResultMedia() {
     return this.status === 'completed' && !this.isNSFW;
   }
 
+  /**
+   * Whether this job produces video output (based on the model used)
+   */
+  get type(): 'image' | 'video' {
+    return this._api.isVideoModelId(this._project.params.modelId) ? 'video' : 'image';
+  }
+
   get enhancedImage() {
     if (!this._enhancementProject) {
       return null;
@@ -199,17 +215,27 @@ class Job extends DataEntity<JobData, JobEventMap> {
 
   /**
    * Get the result URL of the job. This method will make a request to the API to get signed URL.
-   * IMPORTANT: URL expires after 30 minutes, so make sure to download the image as soon as possible.
+   * IMPORTANT: URL expires after 30 minutes, so make sure to download the result as soon as possible.
+   * For video jobs, this returns a video URL. For image jobs, this returns an image URL.
    */
   async getResultUrl(): Promise<string> {
     if (this.data.status !== 'completed') {
       throw new Error('Job is not completed yet');
     }
-    const url = await this._api.downloadUrl({
-      jobId: this.projectId,
-      imageId: this.id,
-      type: 'complete'
-    });
+    let url: string;
+    if (this.type === 'video') {
+      url = await this._api.mediaDownloadUrl({
+        jobId: this.projectId,
+        id: this.id,
+        type: 'complete'
+      });
+    } else {
+      url = await this._api.downloadUrl({
+        jobId: this.projectId,
+        imageId: this.id,
+        type: 'complete'
+      });
+    }
     this._update({ resultUrl: url });
     return url;
   }
@@ -230,6 +256,15 @@ class Job extends DataEntity<JobData, JobEventMap> {
     return this.data.workerName;
   }
 
+  /**
+   * Estimated time remaining in seconds for long-running jobs (e.g., video generation).
+   * Only available for ComfyUI-based workers during inference.
+   * Returns undefined if no ETA has been received.
+   */
+  get etaSeconds() {
+    return this.data.etaSeconds;
+  }
+
   /**
    * Syncs the job data with the data received from the REST API.
    * @internal
@@ -247,11 +282,19 @@ class Job extends DataEntity<JobData, JobEventMap> {
     }
     if (!this.data.resultUrl && delta.status === 'completed' && !data.triggeredNSFWFilter) {
       try {
-        delta.resultUrl = await this._api.downloadUrl({
-          jobId: this.projectId,
-          imageId: this.id,
-          type: 'complete'
-        });
+        if (this.type === 'video') {
+          delta.resultUrl = await this._api.mediaDownloadUrl({
+            jobId: this.projectId,
+            id: this.id,
+            type: 'complete'
+          });
+        } else {
+          delta.resultUrl = await this._api.downloadUrl({
+            jobId: this.projectId,
+            imageId: this.id,
+            type: 'complete'
+          });
+        }
       } catch (error) {
         this._logger.error(error);
       }
@@ -276,8 +319,8 @@ class Job extends DataEntity<JobData, JobEventMap> {
   }
 
   async getResultData() {
-    if (!this.hasResultImage) {
-      throw new Error('No result image available');
+    if (!this.hasResultMedia) {
+      throw new Error('No result media available');
     }
     const url = await this.getResultUrl();
     const response = await fetch(url);
@@ -297,6 +340,10 @@ class Job extends DataEntity<JobData, JobEventMap> {
     strength: EnhancementStrength,
     overrides: { positivePrompt?: string; stylePrompt?: string; tokenType?: TokenType } = {}
   ) {
+    const parentProjectParams = this._project.params;
+    if (parentProjectParams.type !== 'image') {
+      throw new Error('Enhancement is only available for images');
+    }
     if (this.status !== 'completed') {
       throw new Error('Job is not completed yet');
     }
@@ -309,6 +356,7 @@ class Job extends DataEntity<JobData, JobEventMap> {
     }
     const imageData = await this.getResultData();
     const project = await this._api.create({
+      type: 'image',
       ...enhancementDefaults,
       positivePrompt: overrides.positivePrompt || this._project.params.positivePrompt,
       stylePrompt: overrides.stylePrompt || this._project.params.stylePrompt,
@@ -316,7 +364,7 @@ class Job extends DataEntity<JobData, JobEventMap> {
       seed: this.seed || this._project.params.seed,
       startingImage: imageData,
       startingImageStrength: 1 - getEnhacementStrength(strength),
-      sizePreset: this._project.params.sizePreset
+      sizePreset: parentProjectParams.sizePreset
     });
     this._enhancementProject = project;
     this._enhancementProject.on('updated', this.handleEnhancementUpdate);

package/src/Projects/Project.ts

@@ -1,6 +1,6 @@
 import Job, { JobData } from './Job';
 import DataEntity, { EntityEvents } from '../lib/DataEntity';
-import { ProjectParams } from './types';
+import { isImageParams, ProjectParams } from './types';
 import cloneDeep from 'lodash/cloneDeep';
 import ErrorData from '../types/ErrorData';
 import getUUID from '../lib/getUUID';
@@ -93,6 +93,10 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
     return this.data.params;
   }
 
+  get type() {
+    return this.params.type;
+  }
+
   get status() {
     return this.data.status;
   }
@@ -110,10 +114,10 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
    */
   get progress() {
     // Worker can reduce the number of steps in the job, so we need to calculate the progress based on the actual number of steps
-    const stepsPerJob = this.jobs.length ? this.jobs[0].stepCount : this.data.params.steps;
-    const jobCount = this.data.params.numberOfImages;
+    const stepsPerJob = this.jobs.length ? this.jobs[0].stepCount : (this.data.params.steps ?? 0);
+    const jobCount = this.data.params.numberOfMedia;
     const stepsDone = this._jobs.reduce((acc, job) => acc + job.step, 0);
-    return Math.round((stepsDone / (stepsPerJob * jobCount)) * 100);
+    return Math.round((stepsDone / ((stepsPerJob ?? 1) * jobCount)) * 100);
   }
 
   get queuePosition() {
@@ -192,7 +196,7 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
       this._timeout = null;
     }
     if (keys.includes('status') || keys.includes('jobs')) {
-      const allJobsStarted = this.jobs.length >= this.params.numberOfImages;
+      const allJobsStarted = this.jobs.length >= this.params.numberOfMedia;
       const allJobsDone = this.jobs.every((job) => job.finished);
       if (this.data.status === 'completed' && allJobsStarted && allJobsDone) {
         return this.emit('completed', this.resultUrls);
@@ -203,6 +207,16 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
     }
   }
 
+  /**
+   * Refresh the lastUpdated timestamp to prevent timeout.
+   * Used when receiving socket events that indicate the project is still active
+   * (e.g., jobETA events during long-running video generation).
+   * @internal
+   */
+  _keepAlive() {
+    this.lastUpdated = new Date();
+  }
+
   /**
    * This is internal method to add a job to the project. Do not call this directly.
    * @internal
@@ -231,7 +245,11 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
   private _checkForTimeout() {
     if (this.lastUpdated.getTime() + PROJECT_TIMEOUT < Date.now()) {
       this._syncToServer().catch((error) => {
-        this._logger.error(error);
+        // 404 errors are expected when project is still initializing and not yet available via REST API
+        // Only log non-404 errors to avoid confusing users
+        if (error.status !== 404) {
+          this._logger.error(error);
+        }
         this._failedSyncAttempts++;
         if (this._failedSyncAttempts >= MAX_FAILED_SYNC_ATTEMPTS) {
           this._logger.error(
@@ -298,11 +316,13 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
     const delta: Partial<ProjectData> = {
       params: {
         ...this.data.params,
-        numberOfImages: data.imageCount,
-        steps: data.stepCount,
-        numberOfPreviews: data.previewCount
+        numberOfMedia: data.imageCount,
+        steps: data.stepCount
       }
     };
+    if (delta.params && isImageParams(delta.params)) {
+      delta.params.numberOfPreviews = data.previewCount;
+    }
     if (PROJECT_STATUS_MAP[data.status]) {
       delta.status = PROJECT_STATUS_MAP[data.status];
     }

package/src/Projects/createJobRequestMessage.ts

@@ -1,6 +1,54 @@
-import { ProjectParams } from './types';
+import {
+  ImageProjectParams,
+  isImageParams,
+  isVideoParams,
+  ProjectParams,
+  VideoProjectParams
+} from './types';
 import { ControlNetParams, ControlNetParamsRaw } from './types/ControlNetParams';
-import { validateNumber, validateCustomImageSize } from '../lib/validation';
+import {
+  validateNumber,
+  validateCustomImageSize,
+  validateSampler,
+  validateScheduler,
+  validateVideoSize
+} from '../lib/validation';
+import { getVideoWorkflowType, isVideoModel, VIDEO_WORKFLOW_ASSETS } from './utils';
+import { ApiError } from '../ApiClient';
+
+/**
+ * Validate that the provided assets match the workflow requirements.
+ * Throws an error if required assets are missing or forbidden assets are provided.
+ */
+function validateVideoWorkflowAssets(params: VideoProjectParams): void {
+  const workflowType = getVideoWorkflowType(params.modelId);
+  if (!workflowType) return;
+
+  const requirements = VIDEO_WORKFLOW_ASSETS[workflowType];
+  if (!requirements) return;
+  // Check for missing required assets
+  for (const [asset, requirement] of Object.entries(requirements)) {
+    const assetKey = asset as keyof VideoProjectParams;
+    const hasAsset = !!params[assetKey];
+
+    if (requirement === 'required' && !hasAsset) {
+      throw new ApiError(400, {
+        status: 'error',
+        errorCode: 0,
+        message: `${workflowType} workflow requires ${assetKey}. Please provide this asset.`
+      });
+    }
+
+    if (requirement === 'forbidden' && hasAsset) {
+      throw new ApiError(400, {
+        status: 'error',
+        errorCode: 0,
+        message: `${workflowType} workflow does not support ${assetKey}. Please remove this asset.`
+      });
+    }
+  }
+}
+
 // 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 +73,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,
@@ -112,49 +160,120 @@ function getControlNet(params: ControlNetParams): ControlNetParamsRaw[] {
   return [cn];
 }
 
+function applyImageParams(inputKeyframe: Record<string, any>, params: ImageProjectParams) {
+  const keyFrame: Record<string, any> = {
+    ...inputKeyframe,
+    scheduler: validateSampler(params.sampler),
+    timeStepSpacing: validateScheduler(params.scheduler),
+    sizePreset: params.sizePreset,
+    hasContextImage1: !!params.contextImages?.[0],
+    hasContextImage2: !!params.contextImages?.[1],
+    hasContextImage3: !!params.contextImages?.[2]
+  };
+
+  if (params.startingImage) {
+    keyFrame.hasStartingImage = true;
+    keyFrame.strengthIsEnabled = true;
+    keyFrame.strength = 1 - (Number(params.startingImageStrength) || 0.5);
+  }
+
+  if (params.controlNet) {
+    keyFrame.currentControlNetsJob = getControlNet(params.controlNet);
+  }
+  if (params.sizePreset === 'custom') {
+    keyFrame.width = validateCustomImageSize(params.width);
+    keyFrame.height = validateCustomImageSize(params.height);
+  }
+  return keyFrame;
+}
+
+function applyVideoParams(inputKeyframe: Record<string, any>, params: VideoProjectParams) {
+  if (!isVideoModel(params.modelId)) {
+    throw new ApiError(400, {
+      status: 'error',
+      errorCode: 0,
+      message: 'Video generation is only supported for video models.'
+    });
+  }
+  validateVideoWorkflowAssets(params);
+  const keyFrame: Record<string, any> = { ...inputKeyframe };
+  if (params.referenceImage) {
+    keyFrame.hasReferenceImage = true;
+  }
+  if (params.referenceImageEnd) {
+    keyFrame.hasReferenceImageEnd = true;
+  }
+  if (params.referenceAudio) {
+    keyFrame.hasReferenceAudio = true;
+  }
+  if (params.referenceVideo) {
+    keyFrame.hasReferenceVideo = true;
+  }
+
+  // Video generation parameters
+  if (params.frames !== undefined) {
+    keyFrame.frames = params.frames;
+  }
+  if (params.fps !== undefined) {
+    keyFrame.fps = params.fps;
+  }
+  if (params.shift !== undefined) {
+    keyFrame.shift = params.shift;
+  }
+
+  // Validate and set video dimensions (minimum 480px for Wan 2.2 models)
+  if (params.width && params.height) {
+    keyFrame.width = validateVideoSize(params.width, 'width');
+    keyFrame.height = validateVideoSize(params.height, 'height');
+  }
+
+  return keyFrame;
+}
+
 function createJobRequestMessage(id: string, params: ProjectParams) {
   const template = getTemplate();
+  // Base keyFrame with common params
+  let keyFrame: Record<string, any> = {
+    ...template.keyFrames[0],
+    steps: params.steps,
+    guidanceScale: params.guidance,
+    modelID: params.modelId,
+    negativePrompt: params.negativePrompt,
+    seed: params.seed,
+    positivePrompt: params.positivePrompt,
+    stylePrompt: params.stylePrompt
+  };
+
+  switch (params.type) {
+    case 'image':
+      keyFrame = applyImageParams(keyFrame, params);
+      break;
+    case 'video':
+      keyFrame = applyVideoParams(keyFrame, params);
+      break;
+    default:
+      throw new ApiError(400, {
+        status: 'error',
+        errorCode: 0,
+        message: 'Invalid project type. Must be "image" or "video".'
+      });
+  }
+
   const jobRequest: Record<string, any> = {
     ...template,
-    keyFrames: [
-      {
-        ...template.keyFrames[0],
-        scheduler: params.scheduler || null,
-        timeStepSpacing: params.timeStepSpacing || null,
-        steps: params.steps,
-        guidanceScale: params.guidance,
-        modelID: params.modelId,
-        negativePrompt: params.negativePrompt,
-        seed: params.seed,
-        positivePrompt: params.positivePrompt,
-        stylePrompt: params.stylePrompt,
-        hasStartingImage: !!params.startingImage,
-        hasContextImage1: !!params.contextImages?.[0],
-        hasContextImage2: !!params.contextImages?.[1],
-        strengthIsEnabled: !!params.startingImage,
-        strength: !!params.startingImage
-          ? 1 - (Number(params.startingImageStrength) || 0.5)
-          : undefined,
-        sizePreset: params.sizePreset
-      }
-    ],
-    previews: params.numberOfPreviews || 0,
-    numberOfImages: params.numberOfImages,
+    keyFrames: [keyFrame],
+    previews: isImageParams(params) ? params.numberOfPreviews || 0 : 0,
+    numberOfImages: params.numberOfMedia || 1,
     jobID: id,
     disableSafety: !!params.disableNSFWFilter,
     tokenType: params.tokenType,
-    outputFormat: params.outputFormat || 'png'
+    outputFormat: params.outputFormat || (isVideoParams(params) ? 'mp4' : 'png')
   };
+
   if (params.network) {
     jobRequest.network = params.network;
   }
-  if (params.controlNet) {
-    jobRequest.keyFrames[0].currentControlNetsJob = getControlNet(params.controlNet);
-  }
-  if (params.sizePreset === 'custom') {
-    jobRequest.keyFrames[0].width = validateCustomImageSize(params.width);
-    jobRequest.keyFrames[0].height = validateCustomImageSize(params.height);
-  }
+
   return jobRequest;
 }