feat: add flex pro

Author: Abyscuit

modules/bfl/flex.ts

@@ -0,0 +1,55 @@
+import * as Flex from '.';
+const BFL_API = 'https://api.bfl.ml';
+const FLEX_PRO = `${BFL_API}/v1/flux-pro-1.1`;
+const FLEX_PRO_ULTRA = `${BFL_API}/v1/flux-pro-1.1-ultra`;
+const GET_RESULT = `${BFL_API}/v1/get_result?id=`;
+
+export async function generateImage(prompt: string): Promise<Flex.Response> {
+  const promptReq: Flex.ProPrompt = {
+    prompt: prompt,
+    width: 1024,
+    height: 768,
+    prompt_upsampling: false,
+    seed: null,
+    safety_tolerance: 6,
+    output_format: 'jpeg',
+  };
+  const options = {
+    method: 'POST',
+    body: JSON.stringify(promptReq),
+    headers: new Headers({
+      'Content-Type': 'application/json',
+      'X-Key': `${process.env.BFL_KEY}`,
+    }),
+  };
+  const response: Flex.Response = await fetch(new URL(FLEX_PRO), options).then(
+    data => data.json()
+  );
+
+  console.log('response:\n', response);
+
+  let imgResponse = null;
+  while (imgResponse === null) {
+    await sleep(500);
+    imgResponse = await getImageFromAPI(response.id);
+    console.log('imgResponse:\n', imgResponse);
+  }
+
+  return imgResponse;
+}
+
+async function getImageFromAPI(
+  id: string
+): Promise<Flex.GetResponse | Flex.Error | null> {
+  const res: Flex.GetResponse = await fetch(`${GET_RESULT}${id}`).then(data =>
+    data.json()
+  );
+  console.log('res:', res);
+  if (res.status === 'Ready') return res;
+
+  return null;
+}
+
+async function sleep(msec: number) {
+  return new Promise(resolve => setTimeout(resolve, msec));
+}

modules/bfl/index.d.ts

@@ -0,0 +1,32 @@
+export interface PostResponse {
+  id: string;
+}
+
+export interface GetResponse {
+  id: string;
+  status: string;
+  result: object | null;
+}
+
+export interface Error {
+  detail: [
+    {
+      loc: [number | string];
+      msg: string;
+      type: string;
+    },
+  ];
+  id: string;
+}
+
+export type Response = PostResponse | GetResponse | Error;
+
+export type ProPrompt = {
+  prompt: string;
+  width: number | null;
+  height: number | null;
+  prompt_upsampling: boolean | null;
+  seed: number | null;
+  safety_tolerance: 0 | 1 | 2 | 3 | 4 | 5 | 6;
+  output_format: string;
+};

routes/flex.ts

@@ -0,0 +1,39 @@
+import express, { Request, Response } from 'express';
+import { generateImage } from '../modules/bfl/flex';
+const router = express.Router();
+
+router.get('/', (req: Request, res: Response) => {
+  // Check if model is specified
+  //   if (!req.query.model) {
+  //     res.json({
+  //       error: {
+  //         message: 'Model not specified!',
+  //         type: 'invalid_model_error',
+  //         param: null,
+  //         code: null,
+  //       },
+  //     });
+  //     return;
+  //   }
+
+  // Check if there is a prompt
+  if (!req.query.prompt) {
+    res.json({
+      error: {
+        message: 'Prompt cannot be empty!',
+        type: 'invalid_model_error',
+        param: null,
+        code: null,
+      },
+    });
+    return;
+  }
+
+  // Set max tokens
+  //   const maxTokens = parseInt(req.query.max_tokens?.toString() ?? '3000');
+  //   const model = req.query.model as string;
+  const prompt = req.query.prompt as string;
+  generateImage(prompt).then(data => res.json(data));
+});
+
+export default router;

routes/index.ts

@@ -2,11 +2,13 @@ import express from 'express';
 // import users from './users';
 import gpt from './gpt';
 import authorize from './authorize';
+import flex from './flex';
 
 const router = express.Router();
 
 // router.use('/users', users);
 router.use('/gpt', gpt);
 router.use('/authorize', authorize);
+router.use('/flex', flex);
 
 export default router;