@sogni-ai/sogni-client 3.0.0-alpha.9 → 3.0.1

package/src/Projects/Job.ts

@@ -4,6 +4,24 @@ 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';
+import { TokenType } from '../types/token';
+
+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 +58,9 @@ export interface JobData {
   previewUrl?: string;
   resultUrl?: string | null;
   error?: ErrorData;
+  positivePrompt?: string;
+  negativePrompt?: string;
+  jobIndex?: number;
 }
 
 export interface JobEventMap extends EntityEvents {
@@ -51,6 +72,7 @@ export interface JobEventMap extends EntityEvents {
 export interface JobOptions {
   api: ProjectsApi;
   logger: Logger;
+  project: Project;
 }
 
 class Job extends DataEntity<JobData, JobEventMap> {
@@ -72,14 +94,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,6 +178,25 @@ 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.
@@ -225,6 +270,59 @@ 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, style or token type.
+   */
+  async enhance(
+    strength: EnhancementStrength,
+    overrides: { positivePrompt?: string; stylePrompt?: string; tokenType?: TokenType } = {}
+  ) {
+    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,
+      tokenType: overrides.tokenType || this._project.params.tokenType,
+      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,37 +210,24 @@ 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) => {
@@ -298,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

@@ -139,7 +139,8 @@ function createJobRequestMessage(id: string, params: ProjectParams) {
     previews: params.numberOfPreviews || 0,
     numberOfImages: params.numberOfImages,
     jobID: id,
-    disableSafety: !!params.disableNSFWFilter
+    disableSafety: !!params.disableNSFWFilter,
+    tokenType: params.tokenType
   };
   if (params.network) {
     jobRequest.network = params.network;

package/src/Projects/index.ts

@@ -1,6 +1,7 @@
 import ApiGroup, { ApiConfig } from '../ApiGroup';
 import {
   AvailableModel,
+  EnhancementStrength,
   EstimateRequest,
   ImageUrlParams,
   ProjectParams,
@@ -16,7 +17,7 @@ import {
 } from '../ApiClient/WebSocketClient/events';
 import Project from './Project';
 import createJobRequestMessage from './createJobRequestMessage';
-import { ApiError, ApiReponse } from '../ApiClient';
+import { ApiError, ApiResponse } from '../ApiClient';
 import { EstimationResponse } from './types/EstimationResponse';
 import { JobEvent, ProjectApiEvents, ProjectEvent } from './types/events';
 import getUUID from '../lib/getUUID';
@@ -24,6 +25,9 @@ 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';
+import { TokenType } from '../types/token';
 
 const sizePresetCache = new Cache<SizePreset[]>(10 * 60 * 1000);
 const GARBAGE_COLLECT_TIMEOUT = 30000;
@@ -67,7 +71,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 +119,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 +130,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;
       }
@@ -263,15 +273,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') {
@@ -337,10 +362,10 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
    */
   async create(data: ProjectParams): Promise<Project> {
     const project = new Project({ ...data }, { api: this, logger: this.client.logger });
-    if (data.startingImage) {
+    if (data.startingImage && data.startingImage !== true) {
       await this.uploadGuideImage(project.id, data.startingImage);
     }
-    if (data.controlNet?.image) {
+    if (data.controlNet?.image && data.controlNet.image !== true) {
       await this.uploadCNImage(project.id, data.controlNet.image);
     }
     const request = createJobRequestMessage(project.id, data);
@@ -356,7 +381,7 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
    * @param projectId
    */
   async get(projectId: string) {
-    const { data } = await this.client.rest.get<ApiReponse<{ project: RawProject }>>(
+    const { data } = await this.client.rest.get<ApiResponse<{ project: RawProject }>>(
       `/v1/projects/${projectId}`
     );
     return data.project;
@@ -441,6 +466,7 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
    */
   async estimateCost({
     network,
+    tokenType,
     model,
     imageCount,
     stepCount,
@@ -452,6 +478,7 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     sizePreset
   }: EstimateRequest) {
     const pathParams = [
+      tokenType || 'sogni',
       network,
       model,
       imageCount,
@@ -471,7 +498,7 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
       pathParams.push(width, height);
     }
     const r = await this.client.socket.get<EstimationResponse>(
-      `/api/v1/job/estimate/${pathParams.join('/')}`
+      `/api/v2/job/estimate/${pathParams.join('/')}`
     );
     return {
       token: r.quote.project.costInToken,
@@ -479,13 +506,26 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     };
   }
 
+  async estimateEnhancementCost(strength: EnhancementStrength, tokenType: TokenType = 'sogni') {
+    return this.estimateCost({
+      network: enhancementDefaults.network,
+      tokenType,
+      model: enhancementDefaults.modelId,
+      imageCount: 1,
+      stepCount: enhancementDefaults.steps,
+      previewCount: 0,
+      cnEnabled: false,
+      startingImageStrength: getEnhacementStrength(strength)
+    });
+  }
+
   /**
    * Get upload URL for image
    * @internal
    * @param params
    */
   async uploadUrl(params: ImageUrlParams) {
-    const r = await this.client.rest.get<ApiReponse<{ uploadUrl: string }>>(
+    const r = await this.client.rest.get<ApiResponse<{ uploadUrl: string }>>(
       `/v1/image/uploadUrl`,
       params
     );
@@ -498,7 +538,7 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
    * @param params
    */
   async downloadUrl(params: ImageUrlParams) {
-    const r = await this.client.rest.get<ApiReponse<{ downloadUrl: string }>>(
+    const r = await this.client.rest.get<ApiResponse<{ downloadUrl: string }>>(
       `/v1/image/downloadUrl`,
       params
     );

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
@@ -53,8 +54,13 @@ export interface ControlNetParams {
   name: ControlNetName;
   /**
    * ControlNet input image
+   * Supported types:
+   * `File` - file object from input[type=file]
+   * `Buffer` - Node.js buffer object with image data
+   * `Blob` - blob object with image data
+   * `true` - indicates that the image is already uploaded to the server
    */
-  image?: File | Buffer | Blob;
+  image?: File | Buffer | Blob | boolean;
   /**
    * ControlNet strength 0 to 1. 0 full control to prompt, 1 full control to ControlNet
    */

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;
@@ -102,12 +103,14 @@ export interface ProjectParams {
    */
   numberOfImages: number;
   /**
-   * Generate images based on starting image.
+   * Generate images based on the starting image.
+   * Supported types:
    * `File` - file object from input[type=file]
-   * `Buffer` - buffer object with image data
+   * `Buffer` - Node.js buffer object with image data
    * `Blob` - blob object with image data
+   * `true` - indicates that the image is already uploaded to the server
    */
-  startingImage?: File | Buffer | Blob;
+  startingImage?: File | Buffer | Blob | boolean;
   /**
    * How strong effect of starting image should be. From 0 to 1, default 0.5
    */
@@ -141,6 +144,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 = {
@@ -156,6 +164,10 @@ export interface EstimateRequest {
    * Network to use. Can be 'fast' or 'relaxed'
    */
   network: SupernetType;
+  /**
+   * Token type
+   */
+  tokenType?: TokenType;
   /**
    * Model ID
    */
@@ -195,3 +207,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/index.ts

@@ -1,10 +1,10 @@
 import ApiGroup from '../ApiGroup';
-import { ApiReponse } from '../ApiClient';
+import { ApiResponse } from '../ApiClient';
 import { LeaderboardItem, LeaderboardParams } from './types';
 
 class StatsApi extends ApiGroup {
   async leaderboard(params: LeaderboardParams) {
-    const res = await this.client.rest.get<ApiReponse<LeaderboardItem[]>>(
+    const res = await this.client.rest.get<ApiResponse<LeaderboardItem[]>>(
       '/v1/leaderboard/',
       params
     );

package/src/index.ts

@@ -1,4 +1,3 @@
-import { AbstractProvider, JsonRpcProvider, getDefaultProvider } from 'ethers';
 // Account API
 import AccountApi from './Account';
 import CurrentAccount from './Account/CurrentAccount';
@@ -18,6 +17,7 @@ import { AvailableModel, ProjectParams, Scheduler, TimeStepSpacing } from './Pro
 import StatsApi from './Stats';
 // Base Types
 import ErrorData from './types/ErrorData';
+import { TokenType } from './types/token';
 
 export type {
   AvailableModel,
@@ -29,7 +29,8 @@ export type {
   ProjectStatus,
   Scheduler,
   SupernetType,
-  TimeStepSpacing
+  TimeStepSpacing,
+  TokenType
 };
 
 export { ApiError, CurrentAccount, Job, Project };
@@ -49,6 +50,13 @@ export interface SogniClientConfig {
    * @internal
    */
   socketEndpoint?: string;
+  /**
+   * Disable WebSocket connection. Useful for testing or when WebSocket is not needed.
+   * Note that many may not work without WebSocket connection.
+   * @experimental
+   * @internal
+   */
+  disableSocket?: boolean;
   /**
    * Which network to use after logging in. Can be 'fast' or 'relaxed'
    */
@@ -62,10 +70,6 @@ export interface SogniClientConfig {
    * @default 'warn'
    **/
   logLevel?: LogLevel;
-  /**
-   * If provided, the client will connect to this JSON-RPC endpoint to interact with the blockchain
-   */
-  jsonRpcUrl?: string;
   /**
    * If true, the client will connect to the testnet. Ignored if jsonRpcUrl is provided
    */
@@ -92,7 +96,7 @@ export class SogniClient {
   }
 
   /**
-   * Instance creation may involve async operations, so we use a static method
+   * Create client instance, with default configuration
    * @param config
    */
   static async createInstance(config: SogniClientConfig): Promise<SogniClient> {
@@ -100,21 +104,21 @@ export class SogniClient {
     const socketEndpoint = config.socketEndpoint || 'wss://socket.sogni.ai';
     const network = config.network || 'fast';
     const logger = config.logger || new DefaultLogger(config.logLevel || 'warn');
-    const isTestnet = config.testnet !== undefined ? config.testnet : true;
+    const isTestnet = config.testnet !== undefined ? config.testnet : false;
 
-    const client = new ApiClient(restEndpoint, socketEndpoint, config.appId, network, logger);
-    let provider: AbstractProvider;
-    if ('jsonRpcUrl' in config) {
-      provider = new JsonRpcProvider(config.jsonRpcUrl);
-    } else {
-      provider = getDefaultProvider(isTestnet ? 84532 : 8453);
-    }
-    const chainId = await provider.getNetwork().then((network) => network.chainId);
+    const client = new ApiClient(
+      restEndpoint,
+      socketEndpoint,
+      config.appId,
+      network,
+      logger,
+      config.disableSocket
+    );
     const eip712 = new EIP712Helper({
-      name: 'Sogni-testnet',
+      name: isTestnet ? 'Sogni-testnet' : 'Sogni AI',
       version: '1',
-      chainId: chainId.toString()
+      chainId: isTestnet ? '84532' : '8453'
     });
-    return new SogniClient({ client, provider, eip712 });
+    return new SogniClient({ client, eip712 });
   }
 }

package/src/lib/utils.ts

@@ -15,3 +15,7 @@ export function decodeRefreshToken(token: string) {
     expiresAt: new Date(data.exp * 1000)
   };
 }
+
+export async function delay(ms: number) {
+  return new Promise((resolve) => setTimeout(resolve, ms));
+}

package/src/types/token.ts

@@ -0,0 +1 @@
+export type TokenType = 'sogni' | 'spark';

package/src/version.ts

@@ -1 +1,2 @@
-export const LIB_VERSION = "1.0.0-alpha.4";
+import { version } from '../package.json';
+export const LIB_VERSION = version;