@sogni-ai/sogni-client 1.0.0-alpha.4 → 1.0.0

package/src/Projects/index.ts

@@ -1,12 +1,6 @@
 import ApiGroup, { ApiConfig } from '../ApiGroup';
-import {
-  AvailableModel,
-  EstimateRequest,
-  ImageUrlParams,
-  ProjectParams,
-  SizePreset,
-  SupportedModel
-} from './types';
+import models from './models.json';
+import { AvailableModel, EstimateRequest, ImageUrlParams, ProjectParams } from './types';
 import {
   JobErrorData,
   JobProgressData,
@@ -22,12 +16,8 @@ import { JobEvent, ProjectApiEvents, ProjectEvent } from './types/events';
 import getUUID from '../lib/getUUID';
 import { RawProject } from './types/RawProject';
 import ErrorData from '../types/ErrorData';
-import { SupernetType } from '../ApiClient/WebSocketClient/types';
-import Cache from '../lib/Cache';
 
-const sizePresetCache = new Cache<SizePreset[]>(10 * 60 * 1000);
 const GARBAGE_COLLECT_TIMEOUT = 10000;
-const MODELS_REFRESH_INTERVAL = 1000 * 60 * 60 * 24; // 24 hours
 
 function mapErrorCodes(code: string): number {
   switch (code) {
@@ -49,10 +39,6 @@ function mapErrorCodes(code: string): number {
 class ProjectsApi extends ApiGroup<ProjectApiEvents> {
   private _availableModels: AvailableModel[] = [];
   private projects: Project[] = [];
-  private _supportedModels: { data: SupportedModel[] | null; updatedAt: Date } = {
-    data: null,
-    updatedAt: new Date(0)
-  };
 
   get availableModels() {
     return this._availableModels;
@@ -79,20 +65,14 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     this.emit('availableModels', this._availableModels);
   }
 
-  private async handleSwarmModels(data: SocketEventMap['swarmModels']) {
-    let models: SupportedModel[] = [];
-    try {
-      models = await this.getSupportedModels();
-    } catch (e) {
-      this.client.logger.error(e);
-    }
-    const modelIndex = models.reduce((acc: Record<string, SupportedModel>, model) => {
-      acc[model.id] = model;
+  private handleSwarmModels(data: SocketEventMap['swarmModels']) {
+    const modelIndex = models.reduce((acc: Record<string, any>, model) => {
+      acc[model.modelId] = model;
       return acc;
     }, {});
     this._availableModels = Object.entries(data).map(([id, workerCount]) => ({
       id,
-      name: modelIndex[id]?.name || id.replace(/-/g, ' '),
+      name: modelIndex[id]?.modelShortName || id.replace(/-/g, ' '),
       workerCount
     }));
     this.emit('availableModels', this._availableModels);
@@ -390,32 +370,10 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     stepCount,
     previewCount,
     cnEnabled,
-    startingImageStrength,
-    width,
-    height,
-    sizePreset
+    startingImageStrength
   }: EstimateRequest) {
-    const pathParams = [
-      network,
-      model,
-      imageCount,
-      stepCount,
-      previewCount,
-      cnEnabled ? 1 : 0,
-      startingImageStrength ? 1 - startingImageStrength : 0
-    ];
-    if (sizePreset) {
-      const presets = await this.getSizePresets(network, model);
-      const preset = presets.find((p) => p.id === sizePreset);
-      if (!preset) {
-        throw new Error('Invalid size preset');
-      }
-      pathParams.push(preset.width, preset.height);
-    } else if (width && height) {
-      pathParams.push(width, height);
-    }
     const r = await this.client.socket.get<EstimationResponse>(
-      `/api/v1/job/estimate/${pathParams.join('/')}`
+      `/api/v1/job/estimate/${network}/${model}/${imageCount}/${stepCount}/${previewCount}/${cnEnabled ? 1 : 0}/${startingImageStrength ? 1 - startingImageStrength : 0}`
     );
     return {
       token: r.quote.project.costInToken,
@@ -448,74 +406,6 @@ class ProjectsApi extends ApiGroup<ProjectApiEvents> {
     );
     return r.data.downloadUrl;
   }
-
-  async getSupportedModels(forceRefresh = false) {
-    if (
-      this._supportedModels.data &&
-      !forceRefresh &&
-      Date.now() - this._supportedModels.updatedAt.getTime() < MODELS_REFRESH_INTERVAL
-    ) {
-      return this._supportedModels.data;
-    }
-    const models = await this.client.socket.get<SupportedModel[]>(`/api/v1/models/list`);
-    this._supportedModels = { data: models, updatedAt: new Date() };
-    return models;
-  }
-
-  /**
-   * Get supported size presets for the model and network. Size presets are cached for 10 minutes.
-   *
-   * @example
-   * ```ts
-   * const presets = await client.projects.getSizePresets('fast', 'flux1-schnell-fp8');
-   * console.log(presets);
-   * ```
-   *
-   * @param network - 'fast' or 'relaxed'
-   * @param modelId - model id (e.g. 'flux1-schnell-fp8')
-   * @param forceRefresh - force refresh cache
-   * @returns {Promise<{
-   *   label: string;
-   *   id: string;
-   *   width: number;
-   *   height: number;
-   *   ratio: string;
-   *   aspect: string;
-   * }[]>}
-   */
-  async getSizePresets(network: SupernetType, modelId: string, forceRefresh = false) {
-    const key = `${network}-${modelId}`;
-    const cached = sizePresetCache.read(key);
-    if (cached && !forceRefresh) {
-      return cached;
-    }
-    const data = await this.client.socket.get<SizePreset[]>(
-      `/api/v1/size-presets/network/${network}/model/${modelId}`
-    );
-    sizePresetCache.write(key, data);
-    return data;
-  }
-
-  /**
-   * Get available models and their worker counts. Normally, you would get list once you connect
-   * to the server, but you can also call this method to get the list of available models manually.
-   * @param network
-   */
-  async getAvailableModels(network: SupernetType): Promise<AvailableModel[]> {
-    const workersByModelSid = await this.client.socket.get<Record<string, number>>(
-      `/api/v1/status/network/${network}/models`
-    );
-    const supportedModels = await this.getSupportedModels();
-    return Object.entries(workersByModelSid).map(([sid, workerCount]) => {
-      const SID = Number(sid);
-      const model = supportedModels.find((m) => m.SID === SID);
-      return {
-        id: model?.id || sid,
-        name: model?.name || sid.replace(/-/g, ' '),
-        workerCount
-      };
-    });
-  }
 }
 
 export default ProjectsApi;