From ed13a48f12baa4acf3e2298314a9b4926838c57d Mon Sep 17 00:00:00 2001
From: hyper <hx.hyper@qq.com>
Date: Fri, 25 Oct 2024 14:23:29 +0800
Subject: [PATCH] queue and append

---
 src/index.ts  | 33 ++++++++++++++++++++++++++++-----
 src/prodia.ts |  6 ++++--
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index b6c8a2a..383988b 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,7 +1,7 @@
 #!node
 import { program } from 'commander';
 import * as fs from 'fs/promises';
-import { defaultProdiaParams, fetchProdia } from './prodia';
+import { awaitProdiaJob, defaultProdiaParams, fetchProdia } from './prodia';
 
 program
     .name('prodia-scripts')
@@ -13,9 +13,9 @@ program
     .description('generate an image with sdxl')
     .argument('[config]', 'path to the config json', './sdxl.gen.json')
     .option('--init', 'create a default sdxl.gen.json')
+    .option('--append [where]', 'saves generated url to a text file')
+    .option('--queue [amount]', 'queues jobs without waiting for them')
     .action(async (config, flags) => {
-        console.log(config, flags);
-
         if (flags.init) {
             console.log('Creating sdxl.gen.json...');
             await fs.writeFile(
@@ -29,8 +29,31 @@ program
             const json = JSON.parse(txt);
 
             console.log(`Fetching prodia api...`);
-            const status = await fetchProdia(json);
-            console.log('done:', status);
+            if (flags.queue) {
+                const queue = parseInt(flags.queue);
+                if (isNaN(queue) || queue < 1)
+                    throw new Error(`invalid queue number=${flags.queue}`);
+
+                for (let i = 0; i < queue; i++) {
+                    console.log(`Fetching ${i + 1}/${queue}...`);
+                    const job = await fetchProdia(json);
+                    if (flags.append)
+                        await fs.appendFile(
+                            flags.append,
+                            `https://images.prodia.xyz/${job.job}.png\n`
+                        );
+                }
+                console.log('Done!');
+            } else {
+                const job = await fetchProdia(json);
+                const status = await awaitProdiaJob(job);
+
+                console.log('Done: ', status);
+
+                if (status.imageUrl && flags.append) {
+                    await fs.appendFile(flags.append, `${status.imageUrl}\n`);
+                }
+            }
         }
     });
 
diff --git a/src/prodia.ts b/src/prodia.ts
index c342c79..e0d0757 100644
--- a/src/prodia.ts
+++ b/src/prodia.ts
@@ -49,7 +49,10 @@ export async function fetchProdia(
         }),
     });
 
-    const json = (await resp.json()) as ProdiaResponse;
+    return (await resp.json()) as ProdiaResponse;
+}
+
+export async function awaitProdiaJob(json: ProdiaResponse) {
     if ('job' in json) {
         while (true) {
             const query = await fetch(
@@ -64,7 +67,6 @@ export async function fetchProdia(
 
             const status = (await query.json()) as ProdiaResponse;
             if (status.imageUrl) return status;
-
             console.log(status);
 
             await new Promise((r) => setTimeout(r, 5000));