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

package/src/Account/CurrentAccount.ts

@@ -1,5 +1,5 @@
 import DataEntity from '../lib/DataEntity';
-import { BalanceData } from './types';
+import { Balances } from './types';
 import { SupernetType } from '../ApiClient/WebSocketClient/types';
 /**
  * @inline
@@ -16,7 +16,7 @@ export interface AccountData {
    */
   networkStatus: 'connected' | 'disconnected' | 'connecting' | 'switching';
   network: SupernetType | null;
-  balance: BalanceData;
+  balance: Balances;
   walletAddress?: string;
   username?: string;
   token?: string;
@@ -28,10 +28,20 @@ function getDefaults(): AccountData {
     networkStatus: 'disconnected',
     network: null,
     balance: {
-      credit: '0',
-      debit: '0',
-      net: '0',
-      settled: '0'
+      sogni: {
+        credit: '0',
+        debit: '0',
+        net: '0',
+        settled: '0',
+        unclaimed: '0'
+      },
+      spark: {
+        credit: '0',
+        debit: '0',
+        net: '0',
+        settled: '0',
+        unclaimed: '0'
+      }
     },
     walletAddress: undefined,
     username: undefined,

package/src/Account/index.ts

@@ -1,7 +1,7 @@
 import {
   AccountCreateData,
   AccountCreateParams,
-  BalanceData,
+  Balances,
   LoginData,
   Nonce,
   Reward,
@@ -39,7 +39,7 @@ class AccountApi extends ApiGroup {
     this.client.auth.on('updated', this.handleAuthUpdated.bind(this));
   }
 
-  private handleBalanceUpdate(data: BalanceData) {
+  private handleBalanceUpdate(data: Balances) {
     this.currentAccount._update({ balance: data });
   }
 
@@ -214,8 +214,8 @@ class AccountApi extends ApiGroup {
    * // { net: '100.000000', settled: '100.000000', credit: '0.000000', debit: '0.000000' }
    * ```
    */
-  async refreshBalance(): Promise<BalanceData> {
-    const res = await this.client.rest.get<ApiReponse<BalanceData>>('/v1/account/balance');
+  async refreshBalance(): Promise<Balances> {
+    const res = await this.client.rest.get<ApiReponse<Balances>>('/v3/account/balance');
     this.currentAccount._update({ balance: res.data });
     return res.data;
   }
@@ -236,12 +236,11 @@ class AccountApi extends ApiGroup {
    * @param walletAddress
    */
   async walletBalance(walletAddress: string) {
-    const res = await this.client.rest.get<ApiReponse<{ token: string; ether: string }>>(
-      '/v1/wallet/balance',
-      {
-        walletAddress
-      }
-    );
+    const res = await this.client.rest.get<
+      ApiReponse<{ sogni: string; spark: string; ether: string }>
+    >('/v2/wallet/balance', {
+      walletAddress
+    });
     return res.data;
   }
 
@@ -345,7 +344,7 @@ class AccountApi extends ApiGroup {
    */
   async rewards(): Promise<Reward[]> {
     const r =
-      await this.client.rest.get<ApiReponse<{ rewards: RewardRaw[] }>>('/v2/account/rewards');
+      await this.client.rest.get<ApiReponse<{ rewards: RewardRaw[] }>>('/v3/account/rewards');
 
     return r.data.rewards.map(
       (raw: RewardRaw): Reward => ({
@@ -354,6 +353,7 @@ class AccountApi extends ApiGroup {
         title: raw.title,
         description: raw.description,
         amount: raw.amount,
+        tokenType: raw.tokenType,
         claimed: !!raw.claimed,
         canClaim: !!raw.canClaim,
         lastClaim: new Date(raw.lastClaimTimestamp * 1000),
@@ -369,10 +369,12 @@ class AccountApi extends ApiGroup {
    * Claim rewards by reward IDs.
    * @internal
    * @param rewardIds
+   * @param turnstileToken - Turnstile token, required to claim rewards
    */
-  async claimRewards(rewardIds: string[]): Promise<void> {
+  async claimRewards(rewardIds: string[], turnstileToken?: string): Promise<void> {
     await this.client.rest.post('/v2/account/reward/claim', {
-      claims: rewardIds
+      claims: rewardIds,
+      turnstileToken
     });
   }
 

package/src/Account/types.ts

@@ -1,3 +1,5 @@
+import { TokenType } from '../types/token';
+
 export interface Nonce {
   nonce: string;
 }
@@ -27,8 +29,11 @@ export interface BalanceData {
   credit: string;
   debit: string;
   net: string;
+  unclaimed: string;
 }
 
+export type Balances = Record<TokenType, BalanceData>;
+
 export interface TxHistoryParams {
   status: 'completed';
   address: string;
@@ -82,6 +87,7 @@ export interface RewardRaw {
   title: string;
   description: string;
   amount: string;
+  tokenType: TokenType;
   claimed: number;
   canClaim: number;
   lastClaimTimestamp: number;
@@ -94,6 +100,7 @@ export interface Reward {
   title: string;
   description: string;
   amount: string;
+  tokenType: TokenType;
   claimed: boolean;
   canClaim: boolean;
   lastClaim: Date;

package/src/ApiClient/WebSocketClient/events.ts

@@ -1,11 +1,5 @@
 import { SupernetType } from './types';
-
-export type BalanceData = {
-  settled: string;
-  credit: string;
-  debit: string;
-  net: string;
-};
+import { Balances } from '../../Account/types';
 
 export type JobErrorData = {
   jobID: string;
@@ -38,6 +32,9 @@ export type JobStateData =
       jobID: string;
       imgID: string;
       workerName: string;
+      positivePrompt?: string;
+      negativePrompt?: string;
+      jobIndex?: number;
     }
   | {
       jobID: string;
@@ -58,11 +55,25 @@ export type ServerDisconnectData = {
   reason: string;
 };
 
+export type ToastMessage = {
+  type: 'info' | 'success' | 'warning' | 'error';
+  message: string;
+  // Number of milliseconds to show the toast
+  autoClose: number;
+  stickyID: string;
+};
+
+export type ArtistCancelConfirmation = {
+  didCancel: boolean;
+  error_message: string;
+  jobID: string;
+};
+
 export type SocketEventMap = {
   /**
    * @event WebSocketClient#balanceUpdate - Received balance update
    */
-  balanceUpdate: BalanceData;
+  balanceUpdate: Balances;
   /**
    * @event WebSocketClient#changeNetwork - Default network changed
    */
@@ -95,4 +106,10 @@ export type SocketEventMap = {
    * @event WebSocketClient#disconnected - WebSocket connection was closed
    */
   disconnected: ServerDisconnectData;
+  /**
+   * @event WebSocketClient#toastMessage - Toast message received
+   */
+  toastMessage: ToastMessage;
+
+  artistCancelConfirmation: ArtistCancelConfirmation;
 };

package/src/ApiClient/WebSocketClient/index.ts

@@ -28,8 +28,17 @@ class WebSocketClient extends RestClient<SocketEventMap> {
     logger: Logger
   ) {
     const _baseUrl = new URL(baseUrl);
-    if (_baseUrl.protocol === 'wss:') {
-      _baseUrl.protocol = 'https:';
+    switch (_baseUrl.protocol) {
+      case 'http:':
+      case 'ws:':
+        _baseUrl.protocol = 'http:';
+        break;
+      case 'https:':
+      case 'wss:':
+        _baseUrl.protocol = 'https:';
+        break;
+      default:
+        _baseUrl.protocol = 'https:';
     }
     super(_baseUrl.toString(), auth, logger);
     this.appId = appId;
@@ -51,7 +60,8 @@ class WebSocketClient extends RestClient<SocketEventMap> {
     }
     const userAgent = `Sogni/${PROTOCOL_VERSION} (sogni-client) ${LIB_VERSION}`;
     const url = new URL(this.baseUrl);
-    url.protocol = 'wss:';
+    const isNotSecure = url.protocol === 'http:' || url.protocol === 'ws:';
+    url.protocol = isNotSecure ? 'ws:' : 'wss:';
     url.searchParams.set('appId', this.appId);
     url.searchParams.set('clientName', userAgent);
     url.searchParams.set('clientType', 'artist');

package/src/Projects/Job.ts

@@ -4,6 +4,23 @@ import { RawJob, RawProject } from './types/RawProject';
 import ProjectsApi from './index';
 import { Logger } from '../lib/DefaultLogger';
 import getUUID from '../lib/getUUID';
+import { EnhancementStrength } from './types';
+import Project from './Project';
+import { SupernetType } from '../ApiClient/WebSocketClient/types';
+import { getEnhacementStrength } from './utils';
+
+export const enhancementDefaults = {
+  network: 'fast' as SupernetType,
+  modelId: 'flux1-schnell-fp8',
+  positivePrompt: '',
+  negativePrompt: '',
+  stylePrompt: '',
+  startingImageStrength: 0.5,
+  steps: 5,
+  guidance: 1,
+  numberOfImages: 1,
+  numberOfPreviews: 0
+};
 
 export type JobStatus =
   | 'pending'
@@ -40,6 +57,9 @@ export interface JobData {
   previewUrl?: string;
   resultUrl?: string | null;
   error?: ErrorData;
+  positivePrompt?: string;
+  negativePrompt?: string;
+  jobIndex?: number;
 }
 
 export interface JobEventMap extends EntityEvents {
@@ -51,6 +71,7 @@ export interface JobEventMap extends EntityEvents {
 export interface JobOptions {
   api: ProjectsApi;
   logger: Logger;
+  project: Project;
 }
 
 class Job extends DataEntity<JobData, JobEventMap> {
@@ -72,14 +93,18 @@ class Job extends DataEntity<JobData, JobEventMap> {
 
   private readonly _api: ProjectsApi;
   private readonly _logger: Logger;
+  private readonly _project: Project;
+  private _enhancementProject: Project | null = null;
 
   constructor(data: JobData, options: JobOptions) {
     super(data);
 
     this._api = options.api;
     this._logger = options.logger;
+    this._project = options.project;
 
     this.on('updated', this.handleUpdated.bind(this));
+    this.handleEnhancementUpdate = this.handleEnhancementUpdate.bind(this);
   }
 
   get id() {
@@ -152,12 +177,31 @@ class Job extends DataEntity<JobData, JobEventMap> {
     return this.data.error;
   }
 
+  get hasResultImage() {
+    return this.status === 'completed' && !this.isNSFW;
+  }
+
+  get enhancedImage() {
+    if (!this._enhancementProject) {
+      return null;
+    }
+    const project = this._enhancementProject;
+    const job = project.jobs[0];
+    return {
+      status: project.status,
+      progress: project.progress,
+      result: job?.resultUrl || null,
+      error: project.error,
+      getResultUrl: () => job?.getResultUrl()
+    };
+  }
+
   /**
    * 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.
    */
   async getResultUrl(): Promise<string> {
-    if (this.data.status === 'completed') {
+    if (this.data.status !== 'completed') {
       throw new Error('Job is not completed yet');
     }
     const url = await this._api.downloadUrl({
@@ -225,6 +269,58 @@ class Job extends DataEntity<JobData, JobEventMap> {
       this.emit('failed', this.data.error!);
     }
   }
+
+  private handleEnhancementUpdate() {
+    this.emit('updated', ['enhancedImage']);
+  }
+
+  async getResultData() {
+    if (!this.hasResultImage) {
+      throw new Error('No result image available');
+    }
+    const url = await this.getResultUrl();
+    const response = await fetch(url);
+    if (!response.ok) {
+      throw new Error(`Failed to fetch image: ${response.statusText}`);
+    }
+    return response.blob();
+  }
+
+  /**
+   * Enhance the image using the Flux model. This method will create a new project with the
+   * enhancement parameters and use the result image of the current job as the starting image.
+   * @param strength - how much freedom the model has to change the image.
+   * @param overrides - optional parameters to override original prompt or style.
+   */
+  async enhance(
+    strength: EnhancementStrength,
+    overrides: { positivePrompt?: string; stylePrompt?: string } = {}
+  ) {
+    if (this.status !== 'completed') {
+      throw new Error('Job is not completed yet');
+    }
+    if (this.isNSFW) {
+      throw new Error('Job did not pass NSFW filter');
+    }
+    if (this._enhancementProject) {
+      this._enhancementProject.off('updated', this.handleEnhancementUpdate);
+      this._enhancementProject = null;
+    }
+    const imageData = await this.getResultData();
+    const project = await this._api.create({
+      ...enhancementDefaults,
+      positivePrompt: overrides.positivePrompt || this._project.params.positivePrompt,
+      stylePrompt: overrides.stylePrompt || this._project.params.stylePrompt,
+      seed: this.seed || this._project.params.seed,
+      startingImage: imageData,
+      startingImageStrength: 1 - getEnhacementStrength(strength),
+      sizePreset: this._project.params.sizePreset
+    });
+    this._enhancementProject = project;
+    this._enhancementProject.on('updated', this.handleEnhancementUpdate);
+    const images = await project.waitForCompletion();
+    return images[0];
+  }
 }
 
 export default Job;

package/src/Projects/Project.ts

@@ -1,4 +1,4 @@
-import Job, { JobData, JobStatus } from './Job';
+import Job, { JobData } from './Job';
 import DataEntity, { EntityEvents } from '../lib/DataEntity';
 import { ProjectParams } from './types';
 import cloneDeep from 'lodash/cloneDeep';
@@ -50,6 +50,7 @@ export interface ProjectEventMap extends EntityEvents {
   progress: number;
   completed: string[];
   failed: ErrorData;
+  jobStarted: Job;
   jobCompleted: Job;
   jobFailed: Job;
 }
@@ -209,48 +210,47 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
    */
   _addJob(data: JobData | Job) {
     const job =
-      data instanceof Job ? data : new Job(data, { api: this._api, logger: this._logger });
+      data instanceof Job
+        ? data
+        : new Job(data, { api: this._api, logger: this._logger, project: this });
     this._jobs.push(job);
     job.on('updated', () => {
       this.lastUpdated = new Date();
       this.emit('updated', ['jobs']);
     });
+    this.emit('jobStarted', job);
     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) {
+        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;
+          this.jobs.forEach((job) => {
+            if (!job.finished) {
+              job._update({
+                status: 'failed',
+                error: { code: 0, message: 'Job timed out' }
+              });
+            }
+          });
+          this._update({
+            status: 'failed',
+            error: { code: 0, message: 'Project timed out. Please try again or contact support.' }
+          });
         }
       });
     }
@@ -286,7 +286,11 @@ class Project extends DataEntity<ProjectData, ProjectEventMap> {
     // 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 });
+        const jobInstance = Job.fromRaw(data, job, {
+          api: this._api,
+          logger: this._logger,
+          project: this
+        });
         this._addJob(jobInstance);
       }
     }

package/src/Projects/createJobRequestMessage.ts

@@ -119,7 +119,8 @@ function createJobRequestMessage(id: string, params: ProjectParams) {
     keyFrames: [
       {
         ...template.keyFrames[0],
-        scheduler: params.scheduler,
+        scheduler: params.scheduler || null,
+        timeStepSpacing: params.timeStepSpacing || null,
         steps: params.steps,
         guidanceScale: params.guidance,
         modelID: params.modelId,

package/src/Projects/index.ts

@@ -1,6 +1,7 @@
 import ApiGroup, { ApiConfig } from '../ApiGroup';
 import {
   AvailableModel,
+  EnhancementStrength,
   EstimateRequest,
   ImageUrlParams,
   ProjectParams,
@@ -24,9 +25,11 @@ import { RawProject } from './types/RawProject';
 import ErrorData from '../types/ErrorData';
 import { SupernetType } from '../ApiClient/WebSocketClient/types';
 import Cache from '../lib/Cache';
+import { enhancementDefaults } from './Job';
+import { getEnhacementStrength } from './utils';
 
 const sizePresetCache = new Cache<SizePreset[]>(10 * 60 * 1000);
-const GARBAGE_COLLECT_TIMEOUT = 10000;
+const GARBAGE_COLLECT_TIMEOUT = 30000;
 const MODELS_REFRESH_INTERVAL = 1000 * 60 * 60 * 24; // 24 hours
 
 function mapErrorCodes(code: string): number {
@@ -67,7 +70,7 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     this.client.socket.on('jobProgress', this.handleJobProgress.bind(this));
     this.client.socket.on('jobError', this.handleJobError.bind(this));
     this.client.socket.on('jobResult', this.handleJobResult.bind(this));
-    // Listen to server disconnect event
+    // Listen to the server disconnect event
     this.client.on('disconnected', this.handleServerDisconnected.bind(this));
     // Listen to project and job events and update project and job instances
     this.on('project', this.handleProjectEvent.bind(this));
@@ -115,7 +118,10 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
           type: 'initiating',
           projectId: data.jobID,
           jobId: data.imgID,
-          workerName: data.workerName
+          workerName: data.workerName,
+          positivePrompt: data.positivePrompt,
+          negativePrompt: data.negativePrompt,
+          jobIndex: data.jobIndex
         });
         return;
       case 'jobStarted': {
@@ -123,7 +129,10 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
           type: 'started',
           projectId: data.jobID,
           jobId: data.imgID,
-          workerName: data.workerName
+          workerName: data.workerName,
+          positivePrompt: data.positivePrompt,
+          negativePrompt: data.negativePrompt,
+          jobIndex: data.jobIndex
         });
         return;
       }
@@ -241,7 +250,7 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
         this.client.logger.error(e);
       });
       setTimeout(() => {
-        this.projects = this.projects.filter((p) => p.id !== event.projectId);
+        this.projects = this.projects.filter((p) => !p.finished);
       }, GARBAGE_COLLECT_TIMEOUT);
     }
   }
@@ -263,15 +272,30 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     }
     switch (event.type) {
       case 'initiating':
-        job._update({ status: 'initiating', workerName: event.workerName });
+        // positivePrompt and negativePrompt are only received if a Dynamic Prompt was used for the project creating a different prompt for each job
+        job._update({
+          status: 'initiating',
+          workerName: event.workerName,
+          positivePrompt: event.positivePrompt,
+          negativePrompt: event.negativePrompt,
+          jobIndex: event.jobIndex
+        });
         break;
       case 'started':
-        job._update({ status: 'processing', workerName: event.workerName });
+        // positivePrompt and negativePrompt are only received if a Dynamic Prompt was used for the project creating a different prompt for each job
+        job._update({
+          status: 'processing',
+          workerName: event.workerName,
+          positivePrompt: event.positivePrompt,
+          negativePrompt: event.negativePrompt,
+          jobIndex: event.jobIndex
+        });
         break;
       case 'progress':
         job._update({
           status: 'processing',
-          step: event.step,
+          // Jus in case event comes out of order
+          step: Math.max(event.step, job.step),
           stepCount: event.stepCount
         });
         if (project.status !== 'processing') {
@@ -479,6 +503,18 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     };
   }
 
+  async estimateEnhancementCost(strength: EnhancementStrength) {
+    return this.estimateCost({
+      network: enhancementDefaults.network,
+      model: enhancementDefaults.modelId,
+      imageCount: 1,
+      stepCount: enhancementDefaults.steps,
+      previewCount: 0,
+      cnEnabled: false,
+      startingImageStrength: getEnhacementStrength(strength)
+    });
+  }
+
   /**
    * Get upload URL for image
    * @internal

package/src/Projects/types/ControlNetParams.ts

@@ -16,7 +16,8 @@ export type ControlNetName =
   | 'segmentation'
   | 'shuffle'
   | 'softedge'
-  | 'tile';
+  | 'tile'
+  | 'instantid';
 
 /**
  * Raw ControlNet parameters passed to the API

package/src/Projects/types/events.ts

@@ -29,11 +29,17 @@ export interface JobEventBase {
 export interface JobInitiating extends JobEventBase {
   type: 'initiating';
   workerName: string;
+  positivePrompt?: string;
+  negativePrompt?: string;
+  jobIndex?: number;
 }
 
 export interface JobStarted extends JobEventBase {
   type: 'started';
   workerName: string;
+  positivePrompt?: string;
+  negativePrompt?: string;
+  jobIndex?: number;
 }
 
 export interface JobProgress extends JobEventBase {

package/src/Projects/types/index.ts

@@ -1,5 +1,6 @@
 import { SupernetType } from '../../ApiClient/WebSocketClient/types';
 import { ControlNetParams } from './ControlNetParams';
+import { TokenType } from '../../types/token';
 
 export interface SupportedModel {
   id: string;
@@ -141,6 +142,11 @@ export interface ProjectParams {
    * ControlNet model parameters
    */
   controlNet?: ControlNetParams;
+  /**
+   * Select which tokens to use for the project.
+   * If not specified, the Sogni token will be used.
+   */
+  tokenType?: TokenType;
 }
 
 export type ImageUrlParams = {
@@ -195,3 +201,5 @@ export interface EstimateRequest {
    */
   height?: number;
 }
+
+export type EnhancementStrength = 'light' | 'medium' | 'heavy';

package/src/Projects/utils.ts

@@ -0,0 +1,12 @@
+import { EnhancementStrength } from './types';
+
+export function getEnhacementStrength(strength: EnhancementStrength): number {
+  switch (strength) {
+    case 'light':
+      return 0.15;
+    case 'heavy':
+      return 0.49;
+    default:
+      return 0.35;
+  }
+}

package/src/Stats/types.ts

@@ -4,7 +4,9 @@ export type LeaderboardType =
   | 'renderSecCompleteWorker'
   | 'renderSecCompleteArtist'
   | 'renderTokenCompleteWorker'
+  | 'renderTokenCompleteWorker2'
   | 'renderTokenCompleteArtist'
+  | 'renderTokenCompleteArtist2'
   | 'jobCompleteWorker'
   | 'jobCompleteArtist'
   | 'projectCompleteArtist'