Initial implementation of Stable Diffusion

- Needs to add attachment to discord
- Needs to delete image afterward
- Needs to calculate token usage

controllers/stabilityAi/generation_grpc_pb.js

@@ -0,0 +1,71 @@
+// GENERATED CODE -- DO NOT EDIT!
+
+'use strict';
+var grpc = require('grpc');
+var generation_pb = require('./generation_pb.js');
+var google_protobuf_struct_pb = require('google-protobuf/google/protobuf/struct_pb.js');
+var tensors_pb = require('./tensors_pb.js');
+
+function serialize_gooseai_Answer(arg) {
+  if (!(arg instanceof generation_pb.Answer)) {
+    throw new Error('Expected argument of type gooseai.Answer');
+  }
+  return Buffer.from(arg.serializeBinary());
+}
+
+function deserialize_gooseai_Answer(buffer_arg) {
+  return generation_pb.Answer.deserializeBinary(new Uint8Array(buffer_arg));
+}
+
+function serialize_gooseai_ChainRequest(arg) {
+  if (!(arg instanceof generation_pb.ChainRequest)) {
+    throw new Error('Expected argument of type gooseai.ChainRequest');
+  }
+  return Buffer.from(arg.serializeBinary());
+}
+
+function deserialize_gooseai_ChainRequest(buffer_arg) {
+  return generation_pb.ChainRequest.deserializeBinary(new Uint8Array(buffer_arg));
+}
+
+function serialize_gooseai_Request(arg) {
+  if (!(arg instanceof generation_pb.Request)) {
+    throw new Error('Expected argument of type gooseai.Request');
+  }
+  return Buffer.from(arg.serializeBinary());
+}
+
+function deserialize_gooseai_Request(buffer_arg) {
+  return generation_pb.Request.deserializeBinary(new Uint8Array(buffer_arg));
+}
+
+
+//
+// gRPC services
+//
+var GenerationServiceService = exports.GenerationServiceService = {
+  generate: {
+    path: '/gooseai.GenerationService/Generate',
+    requestStream: false,
+    responseStream: true,
+    requestType: generation_pb.Request,
+    responseType: generation_pb.Answer,
+    requestSerialize: serialize_gooseai_Request,
+    requestDeserialize: deserialize_gooseai_Request,
+    responseSerialize: serialize_gooseai_Answer,
+    responseDeserialize: deserialize_gooseai_Answer,
+  },
+  chainGenerate: {
+    path: '/gooseai.GenerationService/ChainGenerate',
+    requestStream: false,
+    responseStream: true,
+    requestType: generation_pb.ChainRequest,
+    responseType: generation_pb.Answer,
+    requestSerialize: serialize_gooseai_ChainRequest,
+    requestDeserialize: deserialize_gooseai_ChainRequest,
+    responseSerialize: serialize_gooseai_Answer,
+    responseDeserialize: deserialize_gooseai_Answer,
+  },
+};
+
+exports.GenerationServiceClient = grpc.makeGenericClientConstructor(GenerationServiceService);

controllers/stabilityAi/generation_pb_service.js

@@ -0,0 +1,117 @@
+// package: gooseai
+// file: generation.proto
+
+var generation_pb = require("./generation_pb");
+var grpc = require("@improbable-eng/grpc-web").grpc;
+
+var GenerationService = (function () {
+  function GenerationService() {}
+  GenerationService.serviceName = "gooseai.GenerationService";
+  return GenerationService;
+}());
+
+GenerationService.Generate = {
+  methodName: "Generate",
+  service: GenerationService,
+  requestStream: false,
+  responseStream: true,
+  requestType: generation_pb.Request,
+  responseType: generation_pb.Answer
+};
+
+GenerationService.ChainGenerate = {
+  methodName: "ChainGenerate",
+  service: GenerationService,
+  requestStream: false,
+  responseStream: true,
+  requestType: generation_pb.ChainRequest,
+  responseType: generation_pb.Answer
+};
+
+exports.GenerationService = GenerationService;
+
+function GenerationServiceClient(serviceHost, options) {
+  this.serviceHost = serviceHost;
+  this.options = options || {};
+}
+
+GenerationServiceClient.prototype.generate = function generate(requestMessage, metadata) {
+  var listeners = {
+    data: [],
+    end: [],
+    status: []
+  };
+  var client = grpc.invoke(GenerationService.Generate, {
+    request: requestMessage,
+    host: this.serviceHost,
+    metadata: metadata,
+    transport: this.options.transport,
+    debug: this.options.debug,
+    onMessage: function (responseMessage) {
+      listeners.data.forEach(function (handler) {
+        handler(responseMessage);
+      });
+    },
+    onEnd: function (status, statusMessage, trailers) {
+      listeners.status.forEach(function (handler) {
+        handler({ code: status, details: statusMessage, metadata: trailers });
+      });
+      listeners.end.forEach(function (handler) {
+        handler({ code: status, details: statusMessage, metadata: trailers });
+      });
+      listeners = null;
+    }
+  });
+  return {
+    on: function (type, handler) {
+      listeners[type].push(handler);
+      return this;
+    },
+    cancel: function () {
+      listeners = null;
+      client.close();
+    }
+  };
+};
+
+GenerationServiceClient.prototype.chainGenerate = function chainGenerate(requestMessage, metadata) {
+  var listeners = {
+    data: [],
+    end: [],
+    status: []
+  };
+  var client = grpc.invoke(GenerationService.ChainGenerate, {
+    request: requestMessage,
+    host: this.serviceHost,
+    metadata: metadata,
+    transport: this.options.transport,
+    debug: this.options.debug,
+    onMessage: function (responseMessage) {
+      listeners.data.forEach(function (handler) {
+        handler(responseMessage);
+      });
+    },
+    onEnd: function (status, statusMessage, trailers) {
+      listeners.status.forEach(function (handler) {
+        handler({ code: status, details: statusMessage, metadata: trailers });
+      });
+      listeners.end.forEach(function (handler) {
+        handler({ code: status, details: statusMessage, metadata: trailers });
+      });
+      listeners = null;
+    }
+  });
+  return {
+    on: function (type, handler) {
+      listeners[type].push(handler);
+      return this;
+    },
+    cancel: function () {
+      listeners = null;
+      client.close();
+    }
+  };
+};
+
+exports.GenerationServiceClient = GenerationServiceClient;
+

controllers/stabilityController.js

@@ -0,0 +1,60 @@
+// Debug
+const { DebugBuilder } = require("../utilities/debugBuilder");
+const log = new DebugBuilder("server", "stabilityController");
+
+// Modules
+const Generation = require("./stabilityAi/generation_pb");
+const { GenerationServiceClient } = require("./stabilityAi/generation_pb_service");
+
+const { grpc } = require("@improbable-eng/grpc-web");
+const GRPCWeb = grpc;
+
+const { NodeHttpTransport } = require("@improbable-eng/grpc-web-node-http-transport");
+const fs = require("fs");
+const {
+  buildGenerationRequest,
+  executeGenerationRequest,
+  onGenerationComplete,
+} = require("../utilities/stabilityHelpers");
+
+// Set Global Color for this controller
+exports.STABILITY_COLOR = 0xeb34b7;
+
+// This is a NodeJS-specific requirement - browsers implementations should omit this line.
+GRPCWeb.setDefaultTransport(NodeHttpTransport());
+
+// Authenticate using your API key, don't commit your key to a public repository!
+const metadata = new GRPCWeb.Metadata();
+metadata.set("Authorization", "Bearer " + process.env.STABILITY_API_KEY);
+
+// Create a generation client to use with all future requests
+const stabilityClient = new GenerationServiceClient("https://grpc.stability.ai", {});
+
+exports.submitImageGenerationTransaction = async () => {
+    const request = buildGenerationRequest("stable-diffusion-512-v2-1", {
+        type: "text-to-image",
+        prompts: [
+          {
+            text: "A dream of a distant galaxy, by Caspar David Friedrich, matte painting trending on artstation HQ",
+          },
+        ],
+        width: 512,
+        height: 512,
+        samples: 1,
+        cfgScale: 13,
+        steps: 10,
+        sampler: Generation.DiffusionSampler.SAMPLER_K_DPMPP_2M,
+    });
+    log.DEBUG("Stability request: ", request, metadata, stabilityClient);
+    
+    executeGenerationRequest(stabilityClient, request, metadata)
+      .then((response) => {
+        log.DEBUG("Stability Generation response: ". response)
+        onGenerationComplete(response);
+        return;
+      })      
+      .catch((error) => {
+        log.ERROR("Failed to make text-to-image request:", error);
+      }
+    );
+}