使用fastapi和vue替代gradio，实现每2秒更新一次上传状态

README.md

@@ -1,8 +1,17 @@
 # WIP
 # stable-diffusion-webui-baidu-netdisk
 stable-diffusion-webui百度云上传拓展。适用于远程云gpu,colab，jupyterlab等需要保存的场合
+
 <img width="678" alt="image" src="https://user-images.githubusercontent.com/25872019/223519780-8de5919a-341b-4912-bdce-eca859a32927.png">
-暂时只实现同步更新的上传结果，使用bduss登录，持续改进中
+
+使用bduss登录，持续改进中
+
+目前可以实现每2秒更新一次上传状态
+
+<img width="1258" alt="image" src="https://user-images.githubusercontent.com/25872019/223800312-0fa01500-c5de-42da-91d1-cde7a59890ba.png">
+
+
+
 
 暂时需要手动下载对应的[BaiduPCS-Go二进制文件](https://github.com/qjfoidnh/BaiduPCS-Go/releases/tag/v3.9.0)放到拓展根目录。后面会自动下载
 

scripts/setup.py

@@ -1,3 +1,4 @@
+from fastapi import FastAPI, HTTPException
 import modules.scripts as scripts
 import gradio as gr
 import os
@@ -7,14 +8,17 @@ import subprocess
 import platform
 import logging
 import time
-
+import uuid
+import select
+import asyncio
+import subprocess
 from modules import images
 from modules.processing import process_images, Processed
 from modules.processing import Processed
 from modules import script_callbacks, shared
 from modules.shared import opts, cmd_opts, state
 import json
-from typing import Dict, Literal, TypedDict
+from typing import IO, Dict, Literal, TypedDict
 
 cwd = os.path.normpath(os.path.join(__file__, "../../"))
 print(shared.config_filename)
@@ -45,8 +49,8 @@ logger.addHandler(file_handler)
 
 def get_global_conf():
     return {
-        "output_dirs": opts.__getattr__("baidu_netdisk_output_dirs"),
-        "upload_dir": opts.__getattr__("baidu_netdisk_upload_dir"),
+        "output_dirs": opts.data["baidu_netdisk_output_dirs"],
+        "upload_dir": opts.data["baidu_netdisk_upload_dir"],
     }
 
 
@@ -119,14 +123,12 @@ def list_file(cwd="/"):
 
 
 def get_curr_user():
-    # 使用正则表达式解析输出
     match = re.search(
         r"uid:\s*(\d+), 用户名:\s*(\w+),",
         exec_ops("who"),
     )
     if not match:
         return
-    # 获取解析结果
     uid = match.group(1)
     if int(uid) == 0:
         return
@@ -146,37 +148,31 @@ def get_curr_user_name():
 
 not_exists_msg = f"找不到{bin_file_name},下载后放到 {cwd} 文件夹下,重启界面"
 
+
 def upload_file_to_baidu_net_disk(pre_log):
     conf = get_global_conf()
-    dirs = str(conf['output_dirs']).split(',')
-    print(["upload", *dirs, conf['upload_dir']])
-    return exec_ops(["upload", *dirs, conf['upload_dir']])
+    dirs = str(conf["output_dirs"]).split(",")
+    print(["upload", *dirs, conf["upload_dir"]])
+
+    return exec_ops(["upload", *dirs, conf["upload_dir"]])
+
 
 def on_ui_tabs():
     exists = check_bin_exists()
     user = get_curr_user()
     if not exists:
         print(f"\033[31m{not_exists_msg}\033[0m")
-    with gr.Blocks(analytics_enabled=False) as image_browser:
+    with gr.Blocks(analytics_enabled=False) as baidu_netdisk:
         gr.Textbox(not_exists_msg, visible=not exists)
         with gr.Row(visible=bool(exists and not user)) as login_form:
             bduss_input = gr.Textbox(interactive=True, label="输入bduss,完成后回车登录")
         with gr.Row(visible=bool(exists and user)) as operation_form:
-            with gr.Column(scale=5):
-                upload_btn = gr.Button("上传")
+            with gr.Column(scale=2):
                 logout_btn = gr.Button("登出账户")
-            with gr.Column(scale=5):
-                log_text = gr.Textbox(get_curr_user_name(), label="log", elem_id="baidu_netdisk_log")
-
-
-            upload_btn.click(
-                fn=upload_file_to_baidu_net_disk,
-                inputs=log_text,
-                outputs=log_text,
-                # _js="document.querySelector(\"body > gradio-app\").shadowRoot.querySelector(\"#baidu_netdisk_log textarea\").value = 'uploading....'"
-            
-            )
-
+            with gr.Column(scale=8):
+                log_text = gr.HTML(
+                    get_curr_user_name(), elem_id="baidu_netdisk_container"
+                )
             def on_bduss_input_enter(bduss):
                 res = login_by_bduss(bduss=bduss)
                 return (
@@ -196,7 +192,7 @@ def on_ui_tabs():
                 return gr.update(visible=True), gr.update(visible=False)
 
             logout_btn.click(fn=on_logout, outputs=[login_form, operation_form])
-        return ((image_browser, "百度云", "baiduyun"),)
+        return ((baidu_netdisk, "百度云", "baiduyun"),)
 
 
 def on_ui_settings():
@@ -240,6 +236,58 @@ def on_ui_settings():
         )
 
 
+subprocess_cache: dict[str, asyncio.subprocess.Process] = {}
+
+
+def is_io_ready(io: IO[bytes]):
+    return select.select([io], [], [], 0) == ([io], [], [])
+
+
+def baidu_netdisk_api(_: gr.Blocks, app: FastAPI):
+    pre = "/baidu_netdisk/"
+
+    @app.get(f"{pre}hello")
+    async def greeting():
+        return "hello"
+
+    @app.post(f"{pre}upload")
+    async def upload():
+        id = str(uuid.uuid4())
+        conf = get_global_conf()
+        dirs = str(conf["output_dirs"]).split(",")
+        with cd(cwd):
+            process = await asyncio.create_subprocess_exec(
+                bin_file_name,
+                "upload",
+                *dirs,
+                conf["upload_dir"],
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+            )
+
+        subprocess_cache[id] = process
+        return {"id": id}
+
+    @app.get(pre + "upload/status/{id}")
+    async def upload_poll(id):
+        p = subprocess_cache.get(id)
+        if not p:
+            raise HTTPException(status_code=404, detail="找不到该subprocess")
+        running = not isinstance(p.returncode, int)
+        msgs = []
+
+        while True:
+            try:
+                line = await asyncio.wait_for(p.stdout.readline(), timeout=0.3)
+                if not line:
+                    break
+                msgs.append(line)
+            except asyncio.TimeoutError:
+                break
+        return {"running": running, "msgs": msgs, "pCode": p.returncode}
+
+
 script_callbacks.on_ui_settings(on_ui_settings)
 
 script_callbacks.on_ui_tabs(on_ui_tabs)
+script_callbacks.on_app_started(baidu_netdisk_api)

style.css

@@ -0,0 +1 @@
+.action-bar[data-v-4562e979]{margin:16px}.container[data-v-4562e979]{width:100%;height:100%;display:flex;flex-direction:column}.log-list[data-v-4562e979]{flex:1;overflow:auto}#baidu_netdisk_container{max-height:70vh}

vue/.eslintrc.cjs

@@ -0,0 +1,15 @@
+/* eslint-env node */
+require('@rushstack/eslint-patch/modern-module-resolution')
+
+module.exports = {
+  root: true,
+  'extends': [
+    'plugin:vue/vue3-essential',
+    'eslint:recommended',
+    '@vue/eslint-config-typescript',
+    '@vue/eslint-config-prettier/skip-formatting'
+  ],
+  parserOptions: {
+    ecmaVersion: 'latest'
+  }
+}

vue/.prettierrc.json

@@ -0,0 +1,8 @@
+{
+  "$schema": "https://json.schemastore.org/prettierrc",
+  "semi": false,
+  "tabWidth": 2,
+  "singleQuote": true,
+  "printWidth": 100,
+  "trailingComma": "none"
+}

vue/.vscode/extensions.json

@@ -0,0 +1,3 @@
+{
+  "recommendations": ["Vue.volar", "Vue.vscode-typescript-vue-plugin"]
+}

vue/README.md

@@ -0,0 +1,46 @@
+# vue
+
+This template should help get you started developing with Vue 3 in Vite.
+
+## Recommended IDE Setup
+
+[VSCode](https://code.visualstudio.com/) + [Volar](https://marketplace.visualstudio.com/items?itemName=Vue.volar) (and disable Vetur) + [TypeScript Vue Plugin (Volar)](https://marketplace.visualstudio.com/items?itemName=Vue.vscode-typescript-vue-plugin).
+
+## Type Support for `.vue` Imports in TS
+
+TypeScript cannot handle type information for `.vue` imports by default, so we replace the `tsc` CLI with `vue-tsc` for type checking. In editors, we need [TypeScript Vue Plugin (Volar)](https://marketplace.visualstudio.com/items?itemName=Vue.vscode-typescript-vue-plugin) to make the TypeScript language service aware of `.vue` types.
+
+If the standalone TypeScript plugin doesn't feel fast enough to you, Volar has also implemented a [Take Over Mode](https://github.com/johnsoncodehk/volar/discussions/471#discussioncomment-1361669) that is more performant. You can enable it by the following steps:
+
+1. Disable the built-in TypeScript Extension
+    1) Run `Extensions: Show Built-in Extensions` from VSCode's command palette
+    2) Find `TypeScript and JavaScript Language Features`, right click and select `Disable (Workspace)`
+2. Reload the VSCode window by running `Developer: Reload Window` from the command palette.
+
+## Customize configuration
+
+See [Vite Configuration Reference](https://vitejs.dev/config/).
+
+## Project Setup
+
+```sh
+yarn
+```
+
+### Compile and Hot-Reload for Development
+
+```sh
+yarn dev
+```
+
+### Type-Check, Compile and Minify for Production
+
+```sh
+yarn build
+```
+
+### Lint with [ESLint](https://eslint.org/)
+
+```sh
+yarn lint
+```

vue/env.d.ts

@@ -0,0 +1 @@
+/// <reference types="vite/client" />

vue/index.html

@@ -0,0 +1,13 @@
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="UTF-8">
+    <link rel="icon" href="/favicon.ico">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Vite App</title>
+  </head>
+  <body>
+    <div id="zanllp_dev_gradio_fe"></div>
+    <script type="module" src="/src/main.ts"></script>
+  </body>
+</html>

vue/package.json

@@ -0,0 +1,33 @@
+{
+  "name": "vue",
+  "version": "0.0.0",
+  "private": true,
+  "scripts": {
+    "dev": "vite",
+    "build": "run-p type-check build-only",
+    "preview": "vite preview",
+    "build-only": "vite build",
+    "type-check": "vue-tsc --noEmit",
+    "lint": "eslint . --ext .vue,.js,.jsx,.cjs,.mjs,.ts,.tsx,.cts,.mts --fix --ignore-path .gitignore",
+    "format": "prettier --write src/"
+  },
+  "dependencies": {
+    "axios": "^1.3.4",
+    "vue": "^3.2.47"
+  },
+  "devDependencies": {
+    "@rushstack/eslint-patch": "^1.2.0",
+    "@types/node": "^18.14.2",
+    "@vitejs/plugin-vue": "^4.0.0",
+    "@vue/eslint-config-prettier": "^7.1.0",
+    "@vue/eslint-config-typescript": "^11.0.2",
+    "@vue/tsconfig": "^0.1.3",
+    "eslint": "^8.34.0",
+    "eslint-plugin-vue": "^9.9.0",
+    "npm-run-all": "^4.1.5",
+    "prettier": "^2.8.4",
+    "typescript": "~4.8.4",
+    "vite": "^4.1.4",
+    "vue-tsc": "^1.2.0"
+  }
+}

vue/src/App.vue

@@ -0,0 +1,68 @@
+<script setup lang="ts">
+import { onMounted, ref, nextTick } from 'vue'
+import { getUploadTaskStatus, greeting, upload } from './api'
+import { Task } from './util/pollTask'
+const msg = ref('')
+const msgs = ref([] as string[])
+onMounted(async () => {
+  msg.value = await greeting()
+})
+const logListEl = ref<HTMLDivElement>()
+const logListScroll2bottom = async () => {
+  await nextTick()
+  const el = logListEl.value
+  if (el) {
+    el.scrollTop = el.scrollHeight
+  }
+}
+const onUploadBtnClick = async () => {
+  msgs.value = []
+  const { id } = await upload()
+  await Task.run({
+    action: () => getUploadTaskStatus(id),
+    pollInterval: 2000,
+    validator (r) {
+      msgs.value.push(...r.msgs)
+      logListScroll2bottom()
+      return !r.running
+    }
+  }).completedTask
+}
+
+</script>
+
+<template>
+  <div class="container">
+
+    <div class="action-bar">
+      <button class="gr-button gr-button-lg gr-button-secondary" @click="onUploadBtnClick">开始上传</button>
+    </div>
+    <div class="log-list" ref="logListEl">
+      <div v-for="msg, idx in msgs" :key="idx">
+        {{ msg }}
+      </div>
+    </div>
+  </div>
+</template>
+<style scoped>
+.action-bar {
+  margin: 16px;
+}
+.container {
+  width: 100%;
+  height: 100%;
+  display: flex;
+  flex-direction: column;
+}
+
+.log-list {
+  flex: 1;
+  overflow: auto;
+}
+</style>
+
+<style>
+#baidu_netdisk_container {
+  max-height: 70vh;
+}
+</style>

vue/src/api/index.ts

@@ -0,0 +1,24 @@
+import axios from 'axios'
+const axiosInst = axios.create({
+  baseURL: '/baidu_netdisk'
+})
+
+export const greeting = async () => {
+  const resp = await axiosInst.get('hello')
+  return resp.data as string
+}
+
+export const upload = async () => {
+  const resp = await axiosInst.post('upload')
+  return resp.data as {
+    id: string
+  }
+}
+
+export const getUploadTaskStatus = async (id: string) => {
+  const resp = await axiosInst.get(`/upload/status/${id}`)
+  return resp.data as {
+    running: boolean,
+    msgs: string[]
+  }
+}

vue/src/main.ts

@@ -0,0 +1,7 @@
+import { createApp } from 'vue'
+import App from './App.vue'
+import { gradioApp, asyncCheck } from './util/index'
+const getContainer = () => document.querySelector("#zanllp_dev_gradio_fe") || gradioApp()?.querySelector("#baidu_netdisk_container")
+asyncCheck(getContainer, 300, Infinity)
+  .then(el => createApp(App).mount(el!))
+

vue/src/util/index.ts

@@ -0,0 +1,20 @@
+export function gradioApp() {
+  const elems = document.getElementsByTagName('gradio-app')
+  const gradioShadowRoot = elems.length == 0 ? null : elems[0].shadowRoot
+  return gradioShadowRoot ? gradioShadowRoot : document;
+}
+export const asyncCheck = async<T> (getter: () => T, checkSize = 100, timeout = 1000) => {
+  return new Promise<T>(x => {
+      const check = (num = 0) => {
+          const target = getter();
+          if (target !== undefined && target !== null) {
+              x(target)
+          } else if (num > timeout / checkSize) {// 超时
+              x(target)
+          } else {
+              setTimeout(() => check(++num), checkSize);
+          }
+      };
+      check();
+  });
+}

vue/src/util/pollTask.ts

@@ -0,0 +1,104 @@
+// from https://github.com/xiachufang/vue3-ts-util/blob/main/src/task.ts
+interface BaseTask<T> {
+  /**
+   * 任务函数，支持异步
+   */
+  action: () => T | Promise<T>
+  /**
+   * 立即执行
+   */
+  immediately?: boolean
+  id?: number
+  /**
+   * 验证器，action结束后调用，为true时结束当前任务
+   */
+  validator?: (r: T) => boolean
+
+  errorHandleMethod?: 'stop' | 'ignore'
+}
+
+interface PollTask<T> extends BaseTask<T> {
+  /**
+   * 轮询间隔
+   */
+  pollInterval: number
+}
+
+export type TaskParams<T> = Omit<PollTask<T>, 'id'> // 去掉 PollTask 接口里面的id
+export type TaskInst<T> = PollTask<T> & { isFinished: boolean; res?: T }
+export class Task {
+  /**  为true时，发生错误不打印 */
+  static silent = false
+
+  static run<T> (taskParams: TaskParams<T>) {
+    const task: TaskInst<T> = {
+      immediately: true, // 默认立即执行
+      id: -1,
+      isFinished: false,
+      errorHandleMethod: 'ignore', // 默认忽略action运行错误，在下个时刻正常运行
+      ...taskParams
+    }
+    let onReject: (t: T) => void
+    let onResolve: (t: T) => void
+    /**
+     * 完成时的promise，
+     * 主动清理时不resolve
+     * */
+    const completedTask = new Promise<T>((resolve, reject) => {
+      onResolve = resolve // 把resolve提取出来，可以减少一层嵌套
+      onReject = reject
+    })
+    /**
+     * 清理轮询任务，多用于在组件卸载时清理
+     *
+     * @example
+     * const { clearTask } = Task.run({....})
+     * onMounted(clearTask)
+     * */
+    const clearTask = () => {
+      task.isFinished = true // 两个都要有，刚好action在执行，阻止运行下个action
+      clearTimeout(task.id) // action未执行，取消任务
+    }
+    const runAction = async () => {
+      try {
+        // eslint-disable-next-line require-atomic-updates
+        task.res = await task.action()
+        // 没有验证器时认为会手动调用clearTask
+        if (task.validator && task.validator(task.res)) {
+          onResolve(task.res)
+          clearTask()
+        }
+      } catch (error: any) {
+        Task.silent || console.error(error)
+        if (task.errorHandleMethod === 'stop') {
+          clearTask()
+          onReject(error)
+        }
+      }
+    }
+    /**
+     * 改用settimeout是为了避免可能某个网络请求阻塞了，后面的又开始运行这种情况
+     */
+    const asyncRunNextAction = () => {
+      if (task.isFinished) {
+        return
+      }
+      task.id = setTimeout(async () => {
+        await runAction()
+        asyncRunNextAction()
+      }, task.pollInterval) as any
+    }
+    /**
+     * 严格保证runNextAction在runAction后执行
+     */
+    setTimeout(async () => {
+      task.immediately && await runAction()
+      asyncRunNextAction()
+    }, 0)
+    return {
+      task, // task 外部不可变
+      clearTask,
+      completedTask
+    }
+  }
+}

vue/tsconfig.json

@@ -0,0 +1,16 @@
+{
+  "extends": "@vue/tsconfig/tsconfig.web.json",
+  "include": ["env.d.ts", "src/**/*", "src/**/*.vue"],
+  "compilerOptions": {
+    "baseUrl": ".",
+    "paths": {
+      "@/*": ["./src/*"]
+    }
+  },
+
+  "references": [
+    {
+      "path": "./tsconfig.node.json"
+    }
+  ]
+}

vue/tsconfig.node.json

@@ -0,0 +1,9 @@
+{
+  "extends": "@vue/tsconfig/tsconfig.node.json",
+  "include": ["vite.config.*", "vitest.config.*", "cypress.config.*", "playwright.config.*"],
+  "compilerOptions": {
+    "composite": true,
+    "types": ["node"],
+    "rootDir": "src"
+  }
+}

vue/vite.config.ts

@@ -0,0 +1,21 @@
+import { fileURLToPath, URL } from 'node:url'
+
+import { defineConfig } from 'vite'
+import vue from '@vitejs/plugin-vue'
+
+// https://vitejs.dev/config/
+export default defineConfig({
+  plugins: [vue()],
+  resolve: {
+    alias: {
+      '@': fileURLToPath(new URL('./src', import.meta.url))
+    }
+  },
+  server: {
+    proxy: {
+      '/baidu_netdisk/': {
+        target: 'http://127.0.0.1:7860/'
+      } 
+    }
+  }
+})