@sogni-ai/sogni-client 0.4.0-aplha.1 → 0.4.0-aplha.11

package/src/Account/index.ts

@@ -259,24 +259,28 @@ class AccountApi extends ApiGroup {
 
   /**
    * Switch between fast and relaxed networks.
-   * Note: This method will close the current WebSocket connection and establish a new one.
-   * Do not call this method if you have any active projects.
+   * This will change default network used to process projects. After switching, you will updated
+   * list of AI models available for on selected network.
    *
    * @example Switch to the fast network
    * ```typescript
-   * client.apiClient.once('connected', ({ network }) => {
-   *  console.log('Switched to the network:', network);
-   * });
    * await client.account.switchNetwork('fast');
+   * console.log('Switched to the fast network, now lets wait until we get list of models');
+   * await client.projects.waitForModels();
    * ```
-   * @param network
+   * @param network - Network type to switch to
    */
-  async switchNetwork(network: SupernetType) {
+  async switchNetwork(network: SupernetType): Promise<SupernetType> {
     this.currentAccount._update({
-      networkStatus: 'connecting',
+      networkStatus: 'switching',
       network: null
     });
-    this.client.socket.switchNetwork(network);
+    const newNetwork = await this.client.socket.switchNetwork(network);
+    this.currentAccount._update({
+      networkStatus: 'connected',
+      network: newNetwork
+    });
+    return newNetwork;
   }
 
   /**

package/src/ApiClient/WebSocketClient/events.ts

@@ -12,7 +12,7 @@ export type JobErrorData = {
   imgID?: string;
   isFromWorker: boolean;
   error_message: string;
-  error: number;
+  error: number | string;
 };
 
 export type JobProgressData = {
@@ -63,6 +63,10 @@ export type SocketEventMap = {
    * @event WebSocketClient#balanceUpdate - Received balance update
    */
   balanceUpdate: BalanceData;
+  /**
+   * @event WebSocketClient#changeNetwork - Default network changed
+   */
+  changeNetwork: { network: SupernetType };
   /**
    * @event WebSocketClient#jobError - Job error occurred
    */

package/src/ApiClient/WebSocketClient/index.ts

@@ -19,9 +19,13 @@ class WebSocketClient extends RestClient<SocketEventMap> {
   private _pingInterval: NodeJS.Timeout | null = null;
 
   constructor(baseUrl: string, appId: string, supernetType: SupernetType, logger: Logger) {
-    super(baseUrl, logger);
+    const _baseUrl = new URL(baseUrl);
+    if (_baseUrl.protocol === 'wss:') {
+      _baseUrl.protocol = 'https:';
+    }
+    super(_baseUrl.toString(), logger);
     this.appId = appId;
-    this.baseUrl = baseUrl;
+    this.baseUrl = _baseUrl.toString();
     this._supernetType = supernetType;
   }
 
@@ -56,6 +60,7 @@ class WebSocketClient extends RestClient<SocketEventMap> {
     }
     const userAgent = `Sogni/${LIB_VERSION} (sogni-client)`;
     const url = new URL(this.baseUrl);
+    url.protocol = 'wss:';
     url.searchParams.set('appId', this.appId);
     url.searchParams.set('clientName', userAgent);
     url.searchParams.set('clientType', 'artist');
@@ -108,10 +113,14 @@ class WebSocketClient extends RestClient<SocketEventMap> {
     }
   }
 
-  switchNetwork(supernetType: SupernetType) {
-    this._supernetType = supernetType;
-    this.disconnect();
-    this.connect();
+  switchNetwork(supernetType: SupernetType): Promise<SupernetType> {
+    return new Promise<SupernetType>(async (resolve, reject) => {
+      this.once('changeNetwork', ({ network }) => {
+        this._supernetType = network;
+        resolve(network);
+      });
+      await this.send('changeNetwork', supernetType);
+    });
   }
 
   /**

package/src/ApiClient/WebSocketClient/messages.ts

@@ -1,7 +1,9 @@
 import { JobRequestRaw } from '../../Projects/createJobRequestMessage';
+import { SupernetType } from './types';
 
 export interface SocketMessageMap {
   jobRequest: JobRequestRaw;
+  changeNetwork: SupernetType;
 }
 
 export type MessageType = keyof SocketMessageMap;

package/src/Projects/Job.ts

@@ -1,5 +1,9 @@
 import DataEntity, { EntityEvents } from '../lib/DataEntity';
 import ErrorData from '../types/ErrorData';
+import { RawJob, RawProject } from './types/RawProject';
+import ProjectsApi from './index';
+import { Logger } from '../lib/DefaultLogger';
+import getUUID from '../lib/getUUID';
 
 export type JobStatus =
   | 'pending'
@@ -9,11 +13,23 @@ export type JobStatus =
   | 'failed'
   | 'canceled';
 
+const JOB_STATUS_MAP: Record<RawJob['status'], JobStatus> = {
+  created: 'pending',
+  queued: 'pending',
+  assigned: 'initiating',
+  initiatingModel: 'initiating',
+  jobStarted: 'processing',
+  jobProgress: 'processing',
+  jobCompleted: 'completed',
+  jobError: 'failed'
+};
+
 /**
  * @inline
  */
 export interface JobData {
   id: string;
+  projectId: string;
   status: JobStatus;
   step: number;
   stepCount: number;
@@ -32,9 +48,37 @@ export interface JobEventMap extends EntityEvents {
   failed: ErrorData;
 }
 
+export interface JobOptions {
+  api: ProjectsApi;
+  logger: Logger;
+}
+
 class Job extends DataEntity<JobData, JobEventMap> {
-  constructor(data: JobData) {
+  static fromRaw(rawProject: RawProject, rawJob: RawJob, options: JobOptions) {
+    return new Job(
+      {
+        id: rawJob.imgID || getUUID(),
+        projectId: rawProject.id,
+        status: JOB_STATUS_MAP[rawJob.status],
+        step: rawJob.performedSteps,
+        stepCount: rawProject.stepCount,
+        workerName: rawJob.worker.name,
+        seed: rawJob.seedUsed,
+        isNSFW: rawJob.triggeredNSFWFilter
+      },
+      options
+    );
+  }
+
+  private readonly _api: ProjectsApi;
+  private readonly _logger: Logger;
+
+  constructor(data: JobData, options: JobOptions) {
     super(data);
+
+    this._api = options.api;
+    this._logger = options.logger;
+
     this.on('updated', this.handleUpdated.bind(this));
   }
 
@@ -42,6 +86,10 @@ class Job extends DataEntity<JobData, JobEventMap> {
     return this.data.id;
   }
 
+  get projectId() {
+    return this.data.projectId;
+  }
+
   /**
    * Current status of the job.
    */
@@ -49,6 +97,10 @@ class Job extends DataEntity<JobData, JobEventMap> {
     return this.data.status;
   }
 
+  get finished() {
+    return ['completed', 'failed', 'canceled'].includes(this.status);
+  }
+
   /**
    * Progress of the job in percentage (0-100).
    */
@@ -116,6 +168,35 @@ class Job extends DataEntity<JobData, JobEventMap> {
     return this.data.workerName;
   }
 
+  /**
+   * Syncs the job data with the data received from the REST API.
+   * @internal
+   * @param data
+   */
+  async _syncWithRestData(data: RawJob) {
+    const delta: Partial<JobData> = {
+      step: data.performedSteps,
+      workerName: data.worker.name,
+      seed: data.seedUsed,
+      isNSFW: data.triggeredNSFWFilter
+    };
+    if (JOB_STATUS_MAP[data.status]) {
+      delta.status = JOB_STATUS_MAP[data.status];
+    }
+    if (!this.data.resultUrl && delta.status === 'completed' && !data.triggeredNSFWFilter) {
+      try {
+        delta.resultUrl = await this._api.downloadUrl({
+          jobId: this.projectId,
+          imageId: this.id,
+          type: 'complete'
+        });
+      } catch (error) {
+        this._logger.error(error);
+      }
+    }
+    this._update(delta);
+  }
+
   private handleUpdated(keys: string[]) {
     if (keys.includes('step') || keys.includes('stepCount')) {
       this.emit('progress', this.progress);

package/src/Projects/Project.ts

@@ -1,11 +1,34 @@
-import Job, { JobData } from './Job';
+import Job, { JobData, JobStatus } from './Job';
 import DataEntity, { EntityEvents } from '../lib/DataEntity';
 import { ProjectParams } from './types';
 import cloneDeep from 'lodash/cloneDeep';
 import ErrorData from '../types/ErrorData';
 import getUUID from '../lib/getUUID';
+import { RawJob, RawProject } from './types/RawProject';
+import ProjectsApi from './index';
+import { Logger } from '../lib/DefaultLogger';
 
-export type ProjectStatus = 'pending' | 'queued' | 'processing' | 'completed' | 'failed';
+// If project is not finished and had no updates for 1 minute, force refresh
+const PROJECT_TIMEOUT = 60 * 1000;
+const MAX_FAILED_SYNC_ATTEMPTS = 3;
+
+export type ProjectStatus =
+  | 'pending'
+  | 'queued'
+  | 'processing'
+  | 'completed'
+  | 'failed'
+  | 'canceled';
+
+const PROJECT_STATUS_MAP: Record<RawProject['status'], ProjectStatus> = {
+  pending: 'pending',
+  active: 'queued',
+  assigned: 'processing',
+  progress: 'processing',
+  completed: 'completed',
+  errored: 'failed',
+  cancelled: 'canceled'
+};
 
 /**
  * @inline
@@ -31,11 +54,20 @@ export interface ProjectEventMap extends EntityEvents {
   jobFailed: Job;
 }
 
+export interface ProjectOptions {
+  api: ProjectsApi;
+  logger: Logger;
+}
+
 class Project extends DataEntity<ProjectData, ProjectEventMap> {
   private _jobs: Job[] = [];
   private _lastEmitedProgress = -1;
+  private readonly _api: ProjectsApi;
+  private readonly _logger: Logger;
+  private _timeout: NodeJS.Timeout | null = null;
+  private _failedSyncAttempts = 0;
 
-  constructor(data: ProjectParams) {
+  constructor(data: ProjectParams, options: ProjectOptions) {
     super({
       id: getUUID(),
       startedAt: new Date(),
@@ -44,6 +76,11 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
       status: 'pending'
     });
 
+    this._api = options.api;
+    this._logger = options.logger;
+
+    this._timeout = setInterval(this._checkForTimeout.bind(this), PROJECT_TIMEOUT);
+
     this.on('updated', this.handleUpdated.bind(this));
   }
 
@@ -59,6 +96,10 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
     return this.data.status;
   }
 
+  get finished() {
+    return ['completed', 'failed', 'canceled'].includes(this.status);
+  }
+
   get error() {
     return this.data.error;
   }
@@ -137,6 +178,11 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
       this.emit('progress', progress);
       this._lastEmitedProgress = progress;
     }
+    // If project is finished stop watching for timeout
+    if (this._timeout && this.finished) {
+      clearInterval(this._timeout!);
+      this._timeout = null;
+    }
     if (keys.includes('status') || keys.includes('jobs')) {
       const allJobsDone = this.jobs.every((job) =>
         ['completed', 'failed', 'canceled'].includes(job.status)
@@ -155,21 +201,104 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
    * @internal
    * @param data
    */
-  _addJob(data: JobData) {
-    const job = new Job(data);
+  _addJob(data: JobData | Job) {
+    const job =
+      data instanceof Job ? data : new Job(data, { api: this._api, logger: this._logger });
     this._jobs.push(job);
     job.on('updated', () => {
+      this.lastUpdated = new Date();
       this.emit('updated', ['jobs']);
     });
     job.on('completed', () => {
       this.emit('jobCompleted', job);
+      this._handleJobFinished(job);
     });
     job.on('failed', () => {
       this.emit('jobFailed', job);
+      this._handleJobFinished(job);
     });
     return job;
   }
 
+  private _handleJobFinished(job: Job) {
+    const finalStatus: JobStatus[] = ['completed', 'failed', 'canceled'];
+    const allJobsDone = this.jobs.every((job) => finalStatus.includes(job.status));
+    // If all jobs are done and project is not already failed or completed, update the project status
+    if (allJobsDone && this.status !== 'failed' && this.status !== 'completed') {
+      const allJobsFailed = this.jobs.every((job) => job.status === 'failed');
+      if (allJobsFailed) {
+        this._update({ status: 'failed' });
+      } else {
+        this._update({ status: 'completed' });
+      }
+    }
+  }
+
+  private _checkForTimeout() {
+    if (this.lastUpdated.getTime() + PROJECT_TIMEOUT < Date.now()) {
+      this._syncToServer().catch((error) => {
+        this._logger.error(error);
+        this._failedSyncAttempts++;
+        if (this._failedSyncAttempts > MAX_FAILED_SYNC_ATTEMPTS) {
+          this._logger.error(
+            `Failed to sync project data after ${MAX_FAILED_SYNC_ATTEMPTS} attempts. Stopping further attempts.`
+          );
+          clearInterval(this._timeout!);
+          this._timeout = null;
+        }
+      });
+    }
+  }
+
+  /**
+   * Sync project data with the data received from the REST API.
+   * @internal
+   */
+  async _syncToServer() {
+    const data = await this._api.get(this.id);
+    const jobData = data.completedWorkerJobs.reduce((acc: Record<string, RawJob>, job) => {
+      const jobId = job.imgID || getUUID();
+      acc[jobId] = job;
+      return acc;
+    }, {});
+    for (const job of this._jobs) {
+      const restJob = jobData[job.id];
+      // This should never happen, but just in case we log a warning
+      if (!restJob) {
+        this._logger.warn(`Job with id ${job.id} not found in the REST project data`);
+        return;
+      }
+      try {
+        await job._syncWithRestData(restJob);
+      } catch (error) {
+        this._logger.error(error);
+        this._logger.error(`Failed to sync job ${job.id}`);
+      }
+      delete jobData[job.id];
+    }
+
+    // If there are any jobs left in jobData, it means they are new jobs that are not in the project yet
+    if (Object.keys(jobData).length) {
+      for (const job of Object.values(jobData)) {
+        const jobInstance = Job.fromRaw(data, job, { api: this._api, logger: this._logger });
+        this._addJob(jobInstance);
+      }
+    }
+
+    const delta: Partial<ProjectData> = {
+      params: {
+        ...this.data.params,
+        numberOfImages: data.imageCount,
+        steps: data.stepCount,
+        numberOfPreviews: data.previewCount
+      }
+    };
+    if (PROJECT_STATUS_MAP[data.status]) {
+      delta.status = PROJECT_STATUS_MAP[data.status];
+    }
+    this._update(delta);
+  }
+
   /**
    * Get full project data snapshot. Can be used to serialize the project and store it in a database.
    */

package/src/Projects/createJobRequestMessage.ts

@@ -67,7 +67,7 @@ function getTemplate() {
 
 function createJobRequestMessage(id: string, params: ProjectParams) {
   const template = getTemplate();
-  return {
+  const jobRequest: Record<string, any> = {
     ...template,
     keyFrames: [
       {
@@ -92,6 +92,10 @@ function createJobRequestMessage(id: string, params: ProjectParams) {
     jobID: id,
     disableSafety: !!params.disableNSFWFilter
   };
+  if (params.network) {
+    jobRequest.network = params.network;
+  }
+  return jobRequest;
 }
 
 export type JobRequestRaw = ReturnType<typeof createJobRequestMessage>;

package/src/Projects/index.ts

@@ -14,9 +14,28 @@ import { ApiError, ApiReponse } from '../ApiClient';
 import { EstimationResponse } from './types/EstimationResponse';
 import { JobEvent, ProjectApiEvents, ProjectEvent } from './types/events';
 import getUUID from '../lib/getUUID';
+import { RawProject } from './types/RawProject';
+import ErrorData from '../types/ErrorData';
 
 const GARBAGE_COLLECT_TIMEOUT = 10000;
 
+function mapErrorCodes(code: string): number {
+  switch (code) {
+    case 'serverRestarting':
+      return 5001;
+    case 'workerDisconnected':
+      return 5002;
+    case 'jobTimedOut':
+      return 5003;
+    case 'artistCanceled':
+      return 5004;
+    case 'workerCancelled':
+      return 5005;
+    default:
+      return 5000;
+  }
+}
+
 class ProjectsApi extends ApiGroup<ProjectApiEvents> {
   private _availableModels: AvailableModel[] = [];
   private projects: Project[] = [];
@@ -28,6 +47,7 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
   constructor(config: ApiConfig) {
     super(config);
     // Listen to server events and emit them as project and job events
+    this.client.socket.on('changeNetwork', this.handleChangeNetwork.bind(this));
     this.client.socket.on('swarmModels', this.handleSwarmModels.bind(this));
     this.client.socket.on('jobState', this.handleJobState.bind(this));
     this.client.socket.on('jobProgress', this.handleJobProgress.bind(this));
@@ -40,6 +60,11 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     this.on('job', this.handleJobEvent.bind(this));
   }
 
+  private handleChangeNetwork() {
+    this._availableModels = [];
+    this.emit('availableModels', this._availableModels);
+  }
+
   private handleSwarmModels(data: SocketEventMap['swarmModels']) {
     const modelIndex = models.reduce((acc: Record<string, any>, model) => {
       acc[model.modelId] = model;
@@ -137,14 +162,25 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
   }
 
   private handleJobError(data: JobErrorData) {
+    const errorCode = Number(data.error);
+    let error: ErrorData;
+    if (!isNaN(errorCode)) {
+      error = {
+        code: errorCode,
+        message: data.error_message
+      };
+    } else {
+      error = {
+        code: mapErrorCodes(data.error as string),
+        originalCode: data.error?.toString(),
+        message: data.error_message
+      };
+    }
     if (!data.imgID) {
       this.emit('project', {
         type: 'error',
         projectId: data.jobID,
-        error: {
-          code: Number(data.error),
-          message: data.error_message
-        }
+        error
       });
       return;
     }
@@ -152,10 +188,7 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
       type: 'error',
       projectId: data.jobID,
       jobId: data.imgID,
-      error: {
-        code: Number(data.error),
-        message: data.error_message
-      }
+      error: error
     });
   }
 
@@ -182,7 +215,11 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
           error: event.error
         });
     }
-    if (project.status === 'completed' || project.status === 'failed') {
+    if (project.finished) {
+      // Sync project data with the server and remove it from the list after some time
+      project._syncToServer().catch((e) => {
+        this.client.logger.error(e);
+      });
       setTimeout(() => {
         this.projects = this.projects.filter((p) => p.id !== event.projectId);
       }, GARBAGE_COLLECT_TIMEOUT);
@@ -198,6 +235,7 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     if (!job) {
       job = project._addJob({
         id: event.jobId,
+        projectId: event.projectId,
         status: 'pending',
         step: 0,
         stepCount: project.params.steps
@@ -205,10 +243,10 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     }
     switch (event.type) {
       case 'initiating':
-        job._update({ status: 'initiating' });
+        job._update({ status: 'initiating', workerName: event.workerName });
         break;
       case 'started':
-        job._update({ status: 'processing' });
+        job._update({ status: 'processing', workerName: event.workerName });
         break;
       case 'progress':
         job._update({
@@ -278,7 +316,7 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
    * @param data
    */
   async create(data: ProjectParams): Promise<Project> {
-    const project = new Project({ ...data });
+    const project = new Project({ ...data }, { api: this, logger: this.client.logger });
     if (data.startingImage) {
       await this.uploadGuideImage(project.id, data.startingImage);
     }
@@ -288,6 +326,19 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     return project;
   }
 
+  /**
+   * Get project by id, this API returns project data from the server only if the project is
+   * completed or failed. If the project is still processing, it will throw 404 error.
+   * @internal
+   * @param projectId
+   */
+  async get(projectId: string) {
+    const { data } = await this.client.rest.get<ApiReponse<RawProject>>(
+      `/v1/projects/${projectId}`
+    );
+    return data;
+  }
+
   private async uploadGuideImage(projectId: string, file: File | Buffer | Blob) {
     const imageId = getUUID();
     const presignedUrl = await this.uploadUrl({
@@ -330,7 +381,12 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     };
   }
 
-  private async uploadUrl(params: ImageUrlParams) {
+  /**
+   * Get upload URL for image
+   * @internal
+   * @param params
+   */
+  async uploadUrl(params: ImageUrlParams) {
     const r = await this.client.rest.get<ApiReponse<{ uploadUrl: string }>>(
       `/v1/image/uploadUrl`,
       params
@@ -338,7 +394,12 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     return r.data.uploadUrl;
   }
 
-  private async downloadUrl(params: ImageUrlParams) {
+  /**
+   * Get download URL for image
+   * @internal
+   * @param params
+   */
+  async downloadUrl(params: ImageUrlParams) {
     const r = await this.client.rest.get<ApiReponse<{ downloadUrl: string }>>(
       `/v1/image/downloadUrl`,
       params