General DiffusionPipeline class, readme updated

Author: dakenf

README.md

@@ -12,12 +12,12 @@ npm i @aislamov/diffusers.js
 
 Browser (see examples/react)
 ```js
-const pipe = StableDiffusionPipeline.fromPretrained('aislamov/stable-diffusion-2-1-base-onnx')
+import { DiffusionPipeline } from '@aislamov/diffusers.js'
+
+const pipe = DiffusionPipeline.fromPretrained('aislamov/stable-diffusion-2-1-base-onnx')
 const images = pipe.run({
   prompt: "an astronaut running a horse",
   numInferenceSteps: 30,
-  width: 512,
-  height: 512,
 })
 
 const canvas = document.getElementById('canvas')
@@ -27,12 +27,13 @@ canvas.getContext('2d').putImageData(data, 0, 0);
 
 Node.js (see examples/node)
 ```js
-const pipe = StableDiffusionPipeline.fromPretrained('aislamov/stable-diffusion-2-1-base-onnx')
+import { DiffusionPipeline } from '@aislamov/diffusers.js'
+import { PNG } from 'pngjs'
+
+const pipe = DiffusionPipeline.fromPretrained('aislamov/stable-diffusion-2-1-base-onnx')
 const images = pipe.run({
   prompt: "an astronaut running a horse",
   numInferenceSteps: 30,
-  width: 512,
-  height: 512,
 })
 
 const data = await images[0].mul(255).round().clipByValue(0, 255).transpose(0, 2, 3, 1)
@@ -46,7 +47,7 @@ p.pack().pipe(fs.createWriteStream('output.png')).on('finish', () => {
 
 'aislamov/stable-diffusion-2-1-base-onnx' model is optimized for GPU and will fail to load without CUDA/DML/WebGPU support. Please use 'cpu' revision on a machine without GPU.
 ```js
-const pipe = StableDiffusionPipeline.fromPretrained('aislamov/stable-diffusion-2-1-base-onnx', 'cpu')
+const pipe = DiffusionPipeline.fromPretrained('aislamov/stable-diffusion-2-1-base-onnx', { revision: 'cpu' })
 ```
 
 ## Running examples

examples/node/package-lock.json

@@ -1,5 +1,5 @@
 import minimist from 'minimist';
-import { StableDiffusionPipeline } from '@aislamov/diffusers.js'
+import { DiffusionPipeline } from '@aislamov/diffusers.js'
 import fs from 'fs'
 import { PNG } from 'pngjs'
 import progress from 'cli-progress'
@@ -19,7 +19,7 @@ function parseCommandLineArgs() {
 
 async function main() {
   const args = parseCommandLineArgs();
-  const pipe = await StableDiffusionPipeline.fromPretrained(
+  const pipe = await DiffusionPipeline.fromPretrained(
     args.m,
     {
       revision: args.rev,

examples/node/src/txt2img.mjs

@@ -1,6 +1,7 @@
 import React, { useEffect, useRef, useState } from 'react'
 import './App.css';
 import {
+  DiffusionPipeline,
   ProgressCallback,
   ProgressCallbackPayload,
   setModelCacheDir,
@@ -69,7 +70,7 @@ function App() {
     }
     setModelState('loading')
     try {
-      pipeline.current = await StableDiffusionPipeline.fromPretrained(
+      pipeline.current = await DiffusionPipeline.fromPretrained(
         'aislamov/stable-diffusion-2-1-base-onnx',
         {
           progressCallback

examples/react/src/App.tsx

@@ -1,6 +1,6 @@
 {
   "name": "@aislamov/diffusers.js",
-  "version": "0.9.1",
+  "version": "0.9.2",
   "license": "MIT",
   "author": {
     "name": "Arthur Islamov",

package.json

@@ -6,6 +6,7 @@ import * as ORT from '@aislamov/onnxruntime-web64'
 
 export * from './pipelines/StableDiffusionPipeline'
 export * from './pipelines/StableDiffusionXLPipeline'
+export * from './pipelines/DiffusionPipeline'
 export * from './pipelines/common'
 export * from './hub'
 export { setModelCacheDir } from '@/hub/browser'

src/index-node.ts

@@ -4,6 +4,7 @@ import { setCacheImpl } from '@/hub'
 
 export * from './pipelines/StableDiffusionPipeline'
 export * from './pipelines/StableDiffusionXLPipeline'
+export * from './pipelines/DiffusionPipeline'
 export * from './pipelines/common'
 export * from './hub'
 export { setModelCacheDir } from '@/hub/browser'

src/index.ts

@@ -0,0 +1,25 @@
+import { PretrainedOptions } from '@/pipelines/common'
+import { GetModelFileOptions } from '@/hub/common'
+import { getModelJSON } from '@/hub'
+import { StableDiffusionPipeline } from '@/pipelines/StableDiffusionPipeline'
+import { StableDiffusionXLPipeline } from '@/pipelines/StableDiffusionXLPipeline'
+
+export class DiffusionPipeline {
+  static async fromPretrained (modelRepoOrPath: string, options?: PretrainedOptions) {
+    const opts: GetModelFileOptions = {
+      ...options,
+    }
+
+    const index = await getModelJSON(modelRepoOrPath, 'model_index.json', true, opts)
+
+    switch (index['_class_name']) {
+      case 'StableDiffusionPipeline':
+      case 'OnnxStableDiffusionPipeline':
+        return StableDiffusionPipeline.fromPretrained(modelRepoOrPath, options)
+      case 'StableDiffusionXLPipeline':
+        return StableDiffusionXLPipeline.fromPretrained(modelRepoOrPath, options)
+      default:
+        throw new Error(`Unknown pipeline type ${index['_class_name']}`)
+    }
+  }
+}