处理后端接口

README.md

@@ -1,3 +1,33 @@
+# server
+基于 VUE/NUXT 构建的AI绘图前端 DEMO
+
+对内网开放读写接口, 以便多台服务器协作调度, 由于NUXT默认是多线程的, 因此需要构建一个公共队列
+```JSON
+[
+  {
+    "MODEL": "model-768",
+    "UID": "XXXXXXXXXXX",
+    "STATUS": "awaiting|produce|end",
+  }
+]
+```
+
+
+
+```BASH
+python scripts/txt2img.py --prompt "a little girl professional photograph of an astronaut riding a horse" --ckpt /data/768-v-ema.ckpt --config /data/768-v-ema.yaml --H 768 --W 768 --device cuda --bf16 --n_samples 1 --n_rows 1
+
+
+stderr:   Running command git clone --filter=blob:none --quiet https://github.com/TencentARC/GFPGAN.git /tmp/pip-req-build-1oercg75
+  Running command git rev-parse -q --verify 'sha^8d2447a2d918f8eba5a4a01463fd48e45126a379'
+  Running command git fetch -q https://github.com/TencentARC/GFPGAN.git 8d2447a2d918f8eba5a4a01463fd48e45126a379
+  Running command git checkout -q 8d2447a2d918f8eba5a4a01463fd48e45126a379
+ERROR: Exception:
+Traceback (most recent call last):
+  File "/data/stable-diffusion-webui/venv/lib/python3.10/site-packages/pip/_vendor/urllib3/response.py", line 438, in _error_catcher
+    yield
+```
+
 # Nuxt 3 Minimal Starter
 
 Look at the [Nuxt 3 documentation](https://nuxt.com/docs/getting-started/introduction) to learn more.

pages/index.vue

@@ -1,17 +1,29 @@
 <template lang="pug">
 div(class="mt-[60px] grid grid-cols-1 lg:grid-cols-4 xl:grid-cols-5 text-white bg-[#05020E] h-[calc(100vh-62px)] border-t border-white/10 2xl:border-t-0 mx-auto 2xl:border-x")
   // 左侧信息
-  aside(class="flex flex-col divide-y divide-white/10 pt-6 space-y-6 lg:overflow-y-auto scrollbar-hide")
+  aside.p-4(class="flex flex-col divide-y divide-white/10 pt-6 space-y-6 lg:overflow-y-auto scrollbar-hide")
     div
-      p.font-bold 迅速に開発できる
+      p.font-bold Filter
+      p.text-gray-400 尝试可以应用于您的图像的不同风格样式
+      div.flex.flex-wrap.items-center.justify-between.pt-2
+        div(v-for="item in models.filter(x=>x.use)" class="w-5/16 h-20 bg-gray-500 rounded-lg flex flex-col-reverse")
+          p.p-1 {{item.name}}
+    div
+      p.font-bold Prompt
       p.text-gray-400 あなたのアイデアを素早く実現するためのフレームワークです。
-      textarea.mt-4.rounded-lg.h-32.w-full.px-4.py-2.bg-gray-500.bg-opacity-5.border.border-gray-500.border-opacity-20.text-gray-500(value="Search" type="text" class="focus:outline-none")
+      textarea.mt-4.rounded-lg.h-32.w-full.px-4.py-2.bg-gray-500.bg-opacity-5.border.border-gray-500.border-opacity-20.text-gray-500(
+        v-model="imageCreate.prompt" type="text" class="focus:outline-none"
+      )
     div
       div.flex.items-center.justify-between
         span.font-bold 从图像中删除特征
         label.switch
-          input(type="checkbox" checked)
+          input(type="checkbox" :checked="imageCreate.exclude_on" @change="imageCreate.exclude_on=!imageCreate.exclude_on")
           div.slider.round
+      div(v-show="imageCreate.exclude_on")
+        textarea.mt-4.rounded-lg.h-32.w-full.px-4.py-2.bg-gray-500.bg-opacity-5.border.border-gray-500.border-opacity-20.text-gray-400(
+          v-model="imageCreate.exclude" type="text" class="focus:outline-none"
+        )
       p.text-gray-400 描述您不希望出现在图像中的细节, 例如颜色, 物体或是风景
     div
       p.font-bold 筛选
@@ -29,12 +41,13 @@ div(class="mt-[60px] grid grid-cols-1 lg:grid-cols-4 xl:grid-cols-5 text-white b
               span(class="absolute left-2 top-2.5 text-gray-500")
                 <svg fill="none" height="20" shape-rendering="geometricPrecision" stroke="currentColor" stroke-linecap="round" stroke-linejoin="round" stroke-width="1.5" viewBox="0 0 24 24" width="20"><path d="M11 17.25a6.25 6.25 0 110-12.5 6.25 6.25 0 010 12.5z"></path><path d="M16 16l4.5 4.5"></path></svg>
           div(class="h-32 lg:h-28")
+          // 风格滤镜列表(弹出/悬浮)
           div(class="grid grid-cols-3 sm:grid-cols-4 lg:grid-cols-3 gap-2 p-2 pt-0")
-            template(v-for="item in 18" :key="item")
+            template(v-for="item in models" :key="item.name")
               button.h-24(class="transition-[transform,opacity] origin-center filter-button relative aspect-[6/5] duration-200 border-2 rounded-lg overflow-hidden  active:border-blue-300/50 border-transparent hover:border-high" aria-label="Select filter style: Colorpop" style="transition-delay: 10ms;")
                 span( style="box-sizing: border-box; display: block; overflow: hidden; width: initial; height: initial; background: none; opacity: 1; border: 0px; margin: 0px; padding: 0px; position: absolute; inset: 0px;")
-                  img( alt="Colorpop" src="https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png" decoding="async" data-nimg="fill" style="position: absolute; inset: 0px; box-sizing: border-box; padding: 0px; border: none; margin: auto; display: block; width: 0px; height: 0px; min-width: 100%; max-width: 100%; min-height: 100%; max-height: 100%; object-fit: cover;" sizes="100vw" srcset="https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 640w, https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 750w, https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 828w, https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 1080w, https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 1200w, https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 1920w, https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 2048w, https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 3840w")
+                  img( alt="Colorpop" :src="item.image" decoding="async" data-nimg="fill" style="position: absolute; inset: 0px; box-sizing: border-box; padding: 0px; border: none; margin: auto; display: block; width: 0px; height: 0px; min-width: 100%; max-width: 100%; min-height: 100%; max-height: 100%; object-fit: cover;" sizes="100vw" srcset="https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 640w, https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 750w, https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 828w, https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 1080w, https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 1200w, https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 1920w, https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 2048w, https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png 3840w")
-                div( class="absolute inset-0 bg-gradient-to-t from-black/90 via-black/40 flex flex-col justify-end text-gray-100 text-left text-sm p-1") Colorpop
+                div( class="absolute inset-0 bg-gradient-to-t from-black/90 via-black/40 flex flex-col justify-end text-gray-100 text-left text-sm p-1") {{ item.name }}
     div
       p.font-bold 通过图像生成图像
       p.text-gray-400 上传或绘制图像以用作灵感
@@ -79,28 +92,14 @@ div(class="mt-[60px] grid grid-cols-1 lg:grid-cols-4 xl:grid-cols-5 text-white b
         </div>
       </fieldset>
       div(class="flex flex-col gap-y-8 py-8")
+        // 选择生成尺寸
         fieldset(class="create-fieldset")
           label 图像尺寸
           p 完成图像的宽度×高度.
           div(class="flex flex-row flex-wrap gap-x-2 gap-y-2")
-            div(class="flex flex-row flex-wrap")
+            div(class="flex flex-row flex-wrap" v-for="item in sizes" :key="item.id")
-              input.hidden(type="radio" class="radio-input" id="image-dim-1" checked="")
+              input.hidden.radio-input(type="radio" :id="item.id" checked="")
-              label.border.border-2.border-gray-500.rounded-md.px-4.py-1(for="image-dim-1" class="!text-[11px]" style="width:98px") 512 × 512
+              label.border.border-2.border-gray-500.rounded-md.px-4.py-1.whitespace-nowrap.text-center(:for="item.id" class="!text-[11px]" style="width:98px") {{ item.width }} x {{ item.height }}
-            div(class="flex flex-row flex-wrap")
-              input.hidden(type="radio" class="radio-input" id="image-dim-2")
-              label.border.border-2.border-gray-500.rounded-md.px-4.py-1(for="image-dim-2" class="!text-[11px]" style="width:98px") 1024 × 1024
-            div(class="flex flex-row flex-wrap")
-              input.hidden(type="radio" class="radio-input" id="image-dim-3")
-              label.border.border-2.border-gray-500.rounded-md.px-4.py-1(for="image-dim-3" class="!text-[11px]" style="width:98px") 640 × 384
-            div(class="flex flex-row flex-wrap")
-              input.hidden(type="radio" class="radio-input" id="image-dim-4")
-              label.border.border-2.border-gray-500.rounded-md.px-4.py-1(for="image-dim-4" class="!text-[11px]" style="width:98px") 384 × 640
-            div(class="flex flex-row flex-wrap")
-              input.hidden(type="radio" class="radio-input" id="image-dim-5")
-              label.border.border-2.border-gray-500.rounded-md.px-4.py-1(for="image-dim-5" class="!text-[11px]" style="width:98px") 768 × 512
-            div(class="flex flex-row flex-wrap")
-              input.hidden(type="radio" class="radio-input" id="image-dim-6")
-              label.border.border-2.border-gray-500.rounded-md.px-4.py-1(for="image-dim-6" class="!text-[11px]" style="width:98px") 512 × 768
           div(class="text-sm grey-100 mt-1")
             p Buy a <a target="_blank" style="color:rgb(118 173 255)" href="/pricing">paid plan</a> for any width or height up to 1536px
         fieldset(class="create-fieldset")
@@ -202,6 +201,9 @@ const views = ref({
 })
 
 const imageCreate = ref({
+  prompt: '渲染提示', // 渲染提示
+  exclude: '排除',   // 排除词汇
+  exclude_on: false, // 排除开关
   // 输入:
   // 输入关键词
   // 排除关键词
@@ -220,6 +222,35 @@ const imageCreate = ref({
   // 私有会话
 })
 
+const sizes = ref([
+  { id:'image-dim-1', width:512, height:512 },
+  { id:'image-dim-2', width:768, height:768 },
+  { id:'image-dim-3', width:1024, height:1024 },
+  { id:'image-dim-4', width:640, height:384 },
+  { id:'image-dim-5', width:384, height:640 },
+  { id:'image-dim-6', width:768, height:512 },
+])
+
+const models = ref([
+  { name:'None',      use: 1, image:'https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png' },
+  { name:'Colorpop1', use: 2, image:'https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png' },
+  { name:'Colorpop2', use: 3, image:'https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png' },
+  { name:'Colorpop3', use: 4, image:'https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png' },
+  { name:'Colorpop3', use: 5, image:'https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png' },
+  { name:'Colorpop3', use: 0, image:'https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png' },
+  { name:'Colorpop3', use: 0, image:'https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png' },
+  { name:'Colorpop3', use: 0, image:'https://storage.googleapis.com/pai-marketing/filters/elizaport_style.png' },
+])
+
+const tasks = ref([
+  // 以图生成图, 正面描述词, 反面描述词, 生成数量, 生成质量, 生成尺寸, 生成模型, 生成图片, 随机种子, 生成指导
+  { tid:'sjaksjka0', model:'Colorpop0', image:'', width:768, height:768, number:1, seed:1, quality:1, prompt:'prompt0' },
+  { tid:'sjaksjka1', model:'Colorpop1', image:'', width:768, height:768, number:1, seed:1, quality:1, prompt:'prompt1' },
+  { tid:'sjaksjka2', model:'Colorpop2', image:'', width:768, height:768, number:1, seed:1, quality:1, prompt:'prompt2' },
+  { tid:'sjaksjka3', model:'Colorpop3', image:'', width:768, height:768, number:1, seed:1, quality:1, prompt:'prompt3' },
+])
+
+
 </script>
 
 <style>

server.py

@@ -0,0 +1,388 @@
+import argparse, os
+import cv2
+import torch
+import numpy as np
+from omegaconf import OmegaConf
+from PIL import Image
+from tqdm import tqdm, trange
+from itertools import islice
+from einops import rearrange
+from torchvision.utils import make_grid
+from pytorch_lightning import seed_everything
+from torch import autocast
+from contextlib import nullcontext
+from imwatermark import WatermarkEncoder
+
+from ldm.util import instantiate_from_config
+from ldm.models.diffusion.ddim import DDIMSampler
+from ldm.models.diffusion.plms import PLMSSampler
+from ldm.models.diffusion.dpm_solver import DPMSolverSampler
+
+torch.set_grad_enabled(False)
+
+def chunk(it, size):
+    it = iter(it)
+    return iter(lambda: tuple(islice(it, size)), ())
+
+
+def load_model_from_config(config, ckpt, device=torch.device("cuda"), verbose=False):
+    print(f"Loading model from {ckpt}")
+    pl_sd = torch.load(ckpt, map_location="cpu")
+    if "global_step" in pl_sd:
+        print(f"Global Step: {pl_sd['global_step']}")
+    sd = pl_sd["state_dict"]
+    model = instantiate_from_config(config.model)
+    m, u = model.load_state_dict(sd, strict=False)
+    if len(m) > 0 and verbose:
+        print("missing keys:")
+        print(m)
+    if len(u) > 0 and verbose:
+        print("unexpected keys:")
+        print(u)
+
+    if device == torch.device("cuda"):
+        model.cuda()
+    elif device == torch.device("cpu"):
+        model.cpu()
+        model.cond_stage_model.device = "cpu"
+    else:
+        raise ValueError(f"Incorrect device name. Received: {device}")
+    model.eval()
+    return model
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        nargs="?",
+        default="a professional photograph of an astronaut riding a triceratops",
+        help="the prompt to render"
+    )
+    parser.add_argument(
+        "--outdir",
+        type=str,
+        nargs="?",
+        help="dir to write results to",
+        default="outputs/txt2img-samples"
+    )
+    parser.add_argument(
+        "--steps",
+        type=int,
+        default=50,
+        help="number of ddim sampling steps",
+    )
+    parser.add_argument(
+        "--plms",
+        action='store_true',
+        help="use plms sampling",
+    )
+    parser.add_argument(
+        "--dpm",
+        action='store_true',
+        help="use DPM (2) sampler",
+    )
+    parser.add_argument(
+        "--fixed_code",
+        action='store_true',
+        help="if enabled, uses the same starting code across all samples ",
+    )
+    parser.add_argument(
+        "--ddim_eta",
+        type=float,
+        default=0.0,
+        help="ddim eta (eta=0.0 corresponds to deterministic sampling",
+    )
+    parser.add_argument(
+        "--n_iter",
+        type=int,
+        default=3,
+        help="sample this often",
+    )
+    parser.add_argument(
+        "--H",
+        type=int,
+        default=512,
+        help="image height, in pixel space",
+    )
+    parser.add_argument(
+        "--W",
+        type=int,
+        default=512,
+        help="image width, in pixel space",
+    )
+    parser.add_argument(
+        "--C",
+        type=int,
+        default=4,
+        help="latent channels",
+    )
+    parser.add_argument(
+        "--f",
+        type=int,
+        default=8,
+        help="downsampling factor, most often 8 or 16",
+    )
+    parser.add_argument(
+        "--n_samples",
+        type=int,
+        default=3,
+        help="how many samples to produce for each given prompt. A.k.a batch size",
+    )
+    parser.add_argument(
+        "--n_rows",
+        type=int,
+        default=0,
+        help="rows in the grid (default: n_samples)",
+    )
+    parser.add_argument(
+        "--scale",
+        type=float,
+        default=9.0,
+        help="unconditional guidance scale: eps = eps(x, empty) + scale * (eps(x, cond) - eps(x, empty))",
+    )
+    parser.add_argument(
+        "--from-file",
+        type=str,
+        help="if specified, load prompts from this file, separated by newlines",
+    )
+    parser.add_argument(
+        "--config",
+        type=str,
+        default="configs/stable-diffusion/v2-inference.yaml",
+        help="path to config which constructs model",
+    )
+    parser.add_argument(
+        "--ckpt",
+        type=str,
+        help="path to checkpoint of model",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="the seed (for reproducible sampling)",
+    )
+    parser.add_argument(
+        "--precision",
+        type=str,
+        help="evaluate at this precision",
+        choices=["full", "autocast"],
+        default="autocast"
+    )
+    parser.add_argument(
+        "--repeat",
+        type=int,
+        default=1,
+        help="repeat each prompt in file this often",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        help="Device on which Stable Diffusion will be run",
+        choices=["cpu", "cuda"],
+        default="cpu"
+    )
+    parser.add_argument(
+        "--torchscript",
+        action='store_true',
+        help="Use TorchScript",
+    )
+    parser.add_argument(
+        "--ipex",
+        action='store_true',
+        help="Use Intel® Extension for PyTorch*",
+    )
+    parser.add_argument(
+        "--bf16",
+        action='store_true',
+        help="Use bfloat16",
+    )
+    opt = parser.parse_args()
+    return opt
+
+
+def put_watermark(img, wm_encoder=None):
+    if wm_encoder is not None:
+        img = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
+        img = wm_encoder.encode(img, 'dwtDct')
+        img = Image.fromarray(img[:, :, ::-1])
+    return img
+
+
+def main(opt):
+    seed_everything(opt.seed)
+
+    config = OmegaConf.load(f"{opt.config}")
+    device = torch.device("cuda") if opt.device == "cuda" else torch.device("cpu")
+    model = load_model_from_config(config, f"{opt.ckpt}", device)
+
+    if opt.plms:
+        sampler = PLMSSampler(model, device=device)
+    elif opt.dpm:
+        sampler = DPMSolverSampler(model, device=device)
+    else:
+        sampler = DDIMSampler(model, device=device)
+
+    os.makedirs(opt.outdir, exist_ok=True)
+    outpath = opt.outdir
+
+    print("Creating invisible watermark encoder (see https://github.com/ShieldMnt/invisible-watermark)...")
+    wm = "SDV2"
+    wm_encoder = WatermarkEncoder()
+    wm_encoder.set_watermark('bytes', wm.encode('utf-8'))
+
+    batch_size = opt.n_samples
+    n_rows = opt.n_rows if opt.n_rows > 0 else batch_size
+    if not opt.from_file:
+        prompt = opt.prompt
+        assert prompt is not None
+        data = [batch_size * [prompt]]
+
+    else:
+        print(f"reading prompts from {opt.from_file}")
+        with open(opt.from_file, "r") as f:
+            data = f.read().splitlines()
+            data = [p for p in data for i in range(opt.repeat)]
+            data = list(chunk(data, batch_size))
+
+    sample_path = os.path.join(outpath, "samples")
+    os.makedirs(sample_path, exist_ok=True)
+    sample_count = 0
+    base_count = len(os.listdir(sample_path))
+    grid_count = len(os.listdir(outpath)) - 1
+
+    start_code = None
+    if opt.fixed_code:
+        start_code = torch.randn([opt.n_samples, opt.C, opt.H // opt.f, opt.W // opt.f], device=device)
+
+    if opt.torchscript or opt.ipex:
+        transformer = model.cond_stage_model.model
+        unet = model.model.diffusion_model
+        decoder = model.first_stage_model.decoder
+        additional_context = torch.cpu.amp.autocast() if opt.bf16 else nullcontext()
+        shape = [opt.C, opt.H // opt.f, opt.W // opt.f]
+
+        if opt.bf16 and not opt.torchscript and not opt.ipex:
+            raise ValueError('Bfloat16 is supported only for torchscript+ipex')
+        if opt.bf16 and unet.dtype != torch.bfloat16:
+            raise ValueError("Use configs/stable-diffusion/intel/ configs with bf16 enabled if " +
+                             "you'd like to use bfloat16 with CPU.")
+        if unet.dtype == torch.float16 and device == torch.device("cpu"):
+            raise ValueError("Use configs/stable-diffusion/intel/ configs for your model if you'd like to run it on CPU.")
+
+        if opt.ipex:
+            import intel_extension_for_pytorch as ipex
+            bf16_dtype = torch.bfloat16 if opt.bf16 else None
+            transformer = transformer.to(memory_format=torch.channels_last)
+            transformer = ipex.optimize(transformer, level="O1", inplace=True)
+
+            unet = unet.to(memory_format=torch.channels_last)
+            unet = ipex.optimize(unet, level="O1", auto_kernel_selection=True, inplace=True, dtype=bf16_dtype)
+
+            decoder = decoder.to(memory_format=torch.channels_last)
+            decoder = ipex.optimize(decoder, level="O1", auto_kernel_selection=True, inplace=True, dtype=bf16_dtype)
+
+        if opt.torchscript:
+            with torch.no_grad(), additional_context:
+                # get UNET scripted
+                if unet.use_checkpoint:
+                    raise ValueError("Gradient checkpoint won't work with tracing. " +
+                    "Use configs/stable-diffusion/intel/ configs for your model or disable checkpoint in your config.")
+
+                img_in = torch.ones(2, 4, 96, 96, dtype=torch.float32)
+                t_in = torch.ones(2, dtype=torch.int64)
+                context = torch.ones(2, 77, 1024, dtype=torch.float32)
+                scripted_unet = torch.jit.trace(unet, (img_in, t_in, context))
+                scripted_unet = torch.jit.optimize_for_inference(scripted_unet)
+                print(type(scripted_unet))
+                model.model.scripted_diffusion_model = scripted_unet
+
+                # get Decoder for first stage model scripted
+                samples_ddim = torch.ones(1, 4, 96, 96, dtype=torch.float32)
+                scripted_decoder = torch.jit.trace(decoder, (samples_ddim))
+                scripted_decoder = torch.jit.optimize_for_inference(scripted_decoder)
+                print(type(scripted_decoder))
+                model.first_stage_model.decoder = scripted_decoder
+
+        prompts = data[0]
+        print("Running a forward pass to initialize optimizations")
+        uc = None
+        if opt.scale != 1.0:
+            uc = model.get_learned_conditioning(batch_size * [""])
+        if isinstance(prompts, tuple):
+            prompts = list(prompts)
+
+        with torch.no_grad(), additional_context:
+            for _ in range(3):
+                c = model.get_learned_conditioning(prompts)
+            samples_ddim, _ = sampler.sample(S=5,
+                                             conditioning=c,
+                                             batch_size=batch_size,
+                                             shape=shape,
+                                             verbose=False,
+                                             unconditional_guidance_scale=opt.scale,
+                                             unconditional_conditioning=uc,
+                                             eta=opt.ddim_eta,
+                                             x_T=start_code)
+            print("Running a forward pass for decoder")
+            for _ in range(3):
+                x_samples_ddim = model.decode_first_stage(samples_ddim)
+
+    precision_scope = autocast if opt.precision=="autocast" or opt.bf16 else nullcontext
+    with torch.no_grad(), \
+        precision_scope(opt.device), \
+        model.ema_scope():
+            all_samples = list()
+            for n in trange(opt.n_iter, desc="Sampling"):
+                for prompts in tqdm(data, desc="data"):
+                    uc = None
+                    if opt.scale != 1.0:
+                        uc = model.get_learned_conditioning(batch_size * [""])
+                    if isinstance(prompts, tuple):
+                        prompts = list(prompts)
+                    c = model.get_learned_conditioning(prompts)
+                    shape = [opt.C, opt.H // opt.f, opt.W // opt.f]
+                    samples, _ = sampler.sample(S=opt.steps,
+                                                     conditioning=c,
+                                                     batch_size=opt.n_samples,
+                                                     shape=shape,
+                                                     verbose=False,
+                                                     unconditional_guidance_scale=opt.scale,
+                                                     unconditional_conditioning=uc,
+                                                     eta=opt.ddim_eta,
+                                                     x_T=start_code)
+
+                    x_samples = model.decode_first_stage(samples)
+                    x_samples = torch.clamp((x_samples + 1.0) / 2.0, min=0.0, max=1.0)
+
+                    for x_sample in x_samples:
+                        x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
+                        img = Image.fromarray(x_sample.astype(np.uint8))
+                        img = put_watermark(img, wm_encoder)
+                        img.save(os.path.join(sample_path, f"{base_count:05}.png"))
+                        base_count += 1
+                        sample_count += 1
+
+                    all_samples.append(x_samples)
+
+            # additionally, save as grid
+            grid = torch.stack(all_samples, 0)
+            grid = rearrange(grid, 'n b c h w -> (n b) c h w')
+            grid = make_grid(grid, nrow=n_rows)
+
+            # to image
+            grid = 255. * rearrange(grid, 'c h w -> h w c').cpu().numpy()
+            grid = Image.fromarray(grid.astype(np.uint8))
+            grid = put_watermark(grid, wm_encoder)
+            grid.save(os.path.join(outpath, f'grid-{grid_count:04}.png'))
+            grid_count += 1
+
+    print(f"Your samples are ready and waiting for you here: \n{outpath} \n"
+          f" \nEnjoy.")
+
+
+if __name__ == "__main__":
+    opt = parse_args()
+    main(opt)

server/api/drawing.ts

@@ -0,0 +1,29 @@
+export default defineEventHandler(async event => {
+    event.context.query = getQuery(event)
+
+    // 获取任务列表的状态进度(普通用户只能看到自己的)
+    if (event.node.req.method === 'GET') {
+        let tasks = await useStorage().getItem(`task:${event.context.query.tid}`)
+        return { 'tasks': tasks }
+    }
+
+    // 保存任务列表的状态进度(只允许BOT写入)
+    if (event.node.req.method === 'PUT') {
+        await useStorage().setItem(`task:${event.context.query.tid}`, event.context.query.tid)
+        return { 'message': 'ok' }
+    }
+
+    // 添加任务到队列(生成一个32位随机字符串作为任务ID)
+    if (event.node.req.method === 'POST') {
+        event.context.query.tid = Math.random().toString(36).substring(2, 34)
+        await useStorage().setItem(`task:${event.context.query.tid}`, event.context.query)
+        return { 'message': 'ok' }
+    }
+
+    // 逐删除自己的指定任务(普通用户只能删除自己的)
+    if (event.node.req.method === 'DELETE') {
+        await useStorage().removeItem(`task:${event.context.query.tid}`)
+        return { 'message': 'ok' }
+    }
+
+})