From 35608d209fb4ae0c5a5d3cf6e8da55c6be665a14 Mon Sep 17 00:00:00 2001 From: Qing Date: Sun, 7 May 2023 13:25:46 +0800 Subject: [PATCH] add stable diffusion progress bar --- lama_cleaner/app/package.json | 3 +- lama_cleaner/app/pnpm-lock.yaml | 98 +++++++++++++++---- lama_cleaner/app/src/App.tsx | 2 + .../src/components/SDProgress/SDProgress.scss | 35 +++++++ .../src/components/SDProgress/SDProgress.tsx | 77 +++++++++++++++ lama_cleaner/app/src/styles/_index.scss | 1 + lama_cleaner/server.py | 11 ++- requirements.txt | 3 +- 8 files changed, 204 insertions(+), 26 deletions(-) create mode 100644 lama_cleaner/app/src/components/SDProgress/SDProgress.scss create mode 100644 lama_cleaner/app/src/components/SDProgress/SDProgress.tsx diff --git a/lama_cleaner/app/package.json b/lama_cleaner/app/package.json index aa3ba76..42c454e 100644 --- a/lama_cleaner/app/package.json +++ b/lama_cleaner/app/package.json @@ -12,6 +12,7 @@ "@radix-ui/react-dropdown-menu": "^2.0.4", "@radix-ui/react-icons": "^1.1.1", "@radix-ui/react-popover": "^1.0.0", + "@radix-ui/react-progress": "^1.0.2", "@radix-ui/react-scroll-area": "^1.0.2", "@radix-ui/react-select": "0.1.2-rc.27", "@radix-ui/react-switch": "^0.1.5", @@ -79,7 +80,7 @@ "resolve-url-loader": "^4.0.0", "sass-loader": "^12.3.0", "semver": "^7.3.5", - "socket.io-client": "^4.5.2", + "socket.io-client": "^4.5.4", "source-map-loader": "^3.0.0", "style-loader": "^3.3.1", "tailwindcss": "^3.0.2", diff --git a/lama_cleaner/app/pnpm-lock.yaml b/lama_cleaner/app/pnpm-lock.yaml index 0fc1920..b01c89d 100644 --- a/lama_cleaner/app/pnpm-lock.yaml +++ b/lama_cleaner/app/pnpm-lock.yaml @@ -9,6 +9,7 @@ specifiers: '@radix-ui/react-dropdown-menu': ^2.0.4 '@radix-ui/react-icons': ^1.1.1 '@radix-ui/react-popover': ^1.0.0 + '@radix-ui/react-progress': ^1.0.2 '@radix-ui/react-scroll-area': ^1.0.2 '@radix-ui/react-select': 0.1.2-rc.27 '@radix-ui/react-switch': ^0.1.5 @@ -86,7 +87,7 @@ specifiers: sass: ^1.49.9 sass-loader: ^12.3.0 semver: ^7.3.5 - socket.io-client: ^4.5.2 + socket.io-client: ^4.5.4 source-map-loader: ^3.0.0 style-loader: ^3.3.1 tailwindcss: ^3.0.2 @@ -106,6 +107,7 @@ dependencies: '@radix-ui/react-dropdown-menu': 2.0.4_zsjcj4gvi24ks76nprapl4hsmq '@radix-ui/react-icons': registry.npmmirror.com/@radix-ui/react-icons/1.1.1_react@17.0.2 '@radix-ui/react-popover': registry.npmmirror.com/@radix-ui/react-popover/1.0.3_zsjcj4gvi24ks76nprapl4hsmq + '@radix-ui/react-progress': 1.0.2_sfoxds7t5ydpegc3knd667wn6m '@radix-ui/react-scroll-area': registry.npmmirror.com/@radix-ui/react-scroll-area/1.0.2_sfoxds7t5ydpegc3knd667wn6m '@radix-ui/react-select': registry.npmmirror.com/@radix-ui/react-select/0.1.2-rc.27_zsjcj4gvi24ks76nprapl4hsmq '@radix-ui/react-switch': registry.npmmirror.com/@radix-ui/react-switch/0.1.5_react@17.0.2 @@ -262,6 +264,15 @@ packages: react-dom: registry.npmmirror.com/react-dom/17.0.2_react@17.0.2 dev: false + /@radix-ui/react-compose-refs/0.1.0_react@17.0.2: + resolution: {integrity: sha512-eyclbh+b77k+69Dk72q3694OHrn9B3QsoIRx7ywX341U9RK1ThgQjMFZoPtmZNQTksXHLNEiefR8hGVeFyInGg==} + peerDependencies: + react: ^16.8 || ^17.0 + dependencies: + '@babel/runtime': 7.20.13 + react: registry.npmmirror.com/react/17.0.2 + dev: false + /@radix-ui/react-compose-refs/1.0.0_react@17.0.2: resolution: {integrity: sha512-0KaSv6sx787/hK3eF53iOkiSLwAGlFMx5lotrqD2pTjB18KbybKoEIgkNZTKC60YECDQTKGTRcDBILwZVqVKvA==} peerDependencies: @@ -437,6 +448,28 @@ packages: react-dom: registry.npmmirror.com/react-dom/17.0.2_react@17.0.2 dev: false + /@radix-ui/react-primitive/0.1.4_react@17.0.2: + resolution: {integrity: sha512-6gSl2IidySupIMJFjYnDIkIWRyQdbu/AHK7rbICPani+LW4b0XdxBXc46og/iZvuwW8pjCS8I2SadIerv84xYA==} + peerDependencies: + react: ^16.8 || ^17.0 + dependencies: + '@babel/runtime': 7.20.13 + '@radix-ui/react-slot': 0.1.2_react@17.0.2 + react: registry.npmmirror.com/react/17.0.2 + dev: false + + /@radix-ui/react-primitive/1.0.1_sfoxds7t5ydpegc3knd667wn6m: + resolution: {integrity: sha512-fHbmislWVkZaIdeF6GZxF0A/NH/3BjrGIYj+Ae6eTmTCr7EB0RQAAVEiqsXK6p3/JcRqVSBQoceZroj30Jj3XA==} + peerDependencies: + react: ^16.8 || ^17.0 || ^18.0 + react-dom: ^16.8 || ^17.0 || ^18.0 + dependencies: + '@babel/runtime': 7.20.13 + '@radix-ui/react-slot': 1.0.1_react@17.0.2 + react: registry.npmmirror.com/react/17.0.2 + react-dom: registry.npmmirror.com/react-dom/17.0.2_react@17.0.2 + dev: false + /@radix-ui/react-primitive/1.0.2_sfoxds7t5ydpegc3knd667wn6m: resolution: {integrity: sha512-zY6G5Qq4R8diFPNwtyoLRZBxzu1Z+SXMlfYpChN7Dv8gvmx9X3qhDqiLWvKseKVJMuedFeU/Sa0Sy/Ia+t06Dw==} peerDependencies: @@ -449,6 +482,19 @@ packages: react-dom: registry.npmmirror.com/react-dom/17.0.2_react@17.0.2 dev: false + /@radix-ui/react-progress/1.0.2_sfoxds7t5ydpegc3knd667wn6m: + resolution: {integrity: sha512-c16RVM43ct2koRcMmPw4b47JWFNs89qe5p4Um9dwoPs0yi+d7It1MJ35EpsX+93o31Mqdwe4vQyu0SrHrygdCg==} + peerDependencies: + react: ^16.8 || ^17.0 || ^18.0 + react-dom: ^16.8 || ^17.0 || ^18.0 + dependencies: + '@babel/runtime': 7.20.13 + '@radix-ui/react-context': 1.0.0_react@17.0.2 + '@radix-ui/react-primitive': 1.0.2_sfoxds7t5ydpegc3knd667wn6m + react: registry.npmmirror.com/react/17.0.2 + react-dom: registry.npmmirror.com/react-dom/17.0.2_react@17.0.2 + dev: false + /@radix-ui/react-roving-focus/1.0.3_sfoxds7t5ydpegc3knd667wn6m: resolution: {integrity: sha512-stjCkIoMe6h+1fWtXlA6cRfikdBzCLp3SnVk7c48cv/uy3DTGoXhN76YaOYUJuy3aEDvDIKwKR5KSmvrtPvQPQ==} peerDependencies: @@ -469,6 +515,16 @@ packages: react-dom: registry.npmmirror.com/react-dom/17.0.2_react@17.0.2 dev: false + /@radix-ui/react-slot/0.1.2_react@17.0.2: + resolution: {integrity: sha512-ADkqfL+agEzEguU3yS26jfB50hRrwf7U4VTwAOZEmi/g+ITcBWe12yM46ueS/UCIMI9Py+gFUaAdxgxafFvY2Q==} + peerDependencies: + react: ^16.8 || ^17.0 + dependencies: + '@babel/runtime': 7.20.13 + '@radix-ui/react-compose-refs': 0.1.0_react@17.0.2 + react: registry.npmmirror.com/react/17.0.2 + dev: false + /@radix-ui/react-slot/1.0.1_react@17.0.2: resolution: {integrity: sha512-avutXAFL1ehGvAXtPquu0YK5oz6ctS474iM3vNGQIkswrVhdrS52e3uoMQBzZhNRAIE0jBnUyXWNmSjGHhCFcw==} peerDependencies: @@ -3197,7 +3253,7 @@ packages: name: '@radix-ui/popper' version: 0.1.0 dependencies: - '@babel/runtime': registry.npmmirror.com/@babel/runtime/7.20.13 + '@babel/runtime': 7.20.13 csstype: registry.npmmirror.com/csstype/3.1.1 dev: false @@ -3225,8 +3281,8 @@ packages: peerDependencies: react: ^16.8 || ^17.0 dependencies: - '@babel/runtime': registry.npmmirror.com/@babel/runtime/7.20.13 - '@radix-ui/react-primitive': registry.npmmirror.com/@radix-ui/react-primitive/0.1.4_react@17.0.2 + '@babel/runtime': 7.20.13 + '@radix-ui/react-primitive': 0.1.4_react@17.0.2 react: registry.npmmirror.com/react/17.0.2 dev: false @@ -3239,8 +3295,8 @@ packages: react: ^16.8 || ^17.0 || ^18.0 react-dom: ^16.8 || ^17.0 || ^18.0 dependencies: - '@babel/runtime': registry.npmmirror.com/@babel/runtime/7.20.13 - '@radix-ui/react-primitive': registry.npmmirror.com/@radix-ui/react-primitive/1.0.1_sfoxds7t5ydpegc3knd667wn6m + '@babel/runtime': 7.20.13 + '@radix-ui/react-primitive': 1.0.1_sfoxds7t5ydpegc3knd667wn6m react: registry.npmmirror.com/react/17.0.2 react-dom: registry.npmmirror.com/react-dom/17.0.2_react@17.0.2 dev: false @@ -3272,11 +3328,11 @@ packages: react: ^16.8 || ^17.0 || ^18.0 react-dom: ^16.8 || ^17.0 || ^18.0 dependencies: - '@babel/runtime': registry.npmmirror.com/@babel/runtime/7.20.13 - '@radix-ui/react-compose-refs': registry.npmmirror.com/@radix-ui/react-compose-refs/1.0.0_react@17.0.2 - '@radix-ui/react-context': registry.npmmirror.com/@radix-ui/react-context/1.0.0_react@17.0.2 - '@radix-ui/react-primitive': registry.npmmirror.com/@radix-ui/react-primitive/1.0.1_sfoxds7t5ydpegc3knd667wn6m - '@radix-ui/react-slot': registry.npmmirror.com/@radix-ui/react-slot/1.0.1_react@17.0.2 + '@babel/runtime': 7.20.13 + '@radix-ui/react-compose-refs': 1.0.0_react@17.0.2 + '@radix-ui/react-context': 1.0.0_react@17.0.2 + '@radix-ui/react-primitive': 1.0.1_sfoxds7t5ydpegc3knd667wn6m + '@radix-ui/react-slot': 1.0.1_react@17.0.2 react: registry.npmmirror.com/react/17.0.2 react-dom: registry.npmmirror.com/react-dom/17.0.2_react@17.0.2 dev: false @@ -4042,7 +4098,7 @@ packages: peerDependencies: react: ^16.8 || ^17.0 dependencies: - '@babel/runtime': registry.npmmirror.com/@babel/runtime/7.20.13 + '@babel/runtime': 7.20.13 '@radix-ui/react-use-layout-effect': registry.npmmirror.com/@radix-ui/react-use-layout-effect/0.1.0_react@17.0.2 react: registry.npmmirror.com/react/17.0.2 dev: false @@ -4055,7 +4111,7 @@ packages: peerDependencies: react: ^16.8 || ^17.0 || ^18.0 dependencies: - '@babel/runtime': registry.npmmirror.com/@babel/runtime/7.20.13 + '@babel/runtime': 7.20.13 '@radix-ui/react-use-layout-effect': registry.npmmirror.com/@radix-ui/react-use-layout-effect/0.1.1-rc.18_react@17.0.2 react: registry.npmmirror.com/react/17.0.2 dev: false @@ -4156,7 +4212,7 @@ packages: peerDependencies: react: ^16.8 || ^17.0 || ^18.0 dependencies: - '@babel/runtime': registry.npmmirror.com/@babel/runtime/7.20.13 + '@babel/runtime': 7.20.13 '@radix-ui/react-use-callback-ref': registry.npmmirror.com/@radix-ui/react-use-callback-ref/0.1.1-rc.18_react@17.0.2 react: registry.npmmirror.com/react/17.0.2 dev: false @@ -4169,7 +4225,7 @@ packages: peerDependencies: react: ^16.8 || ^17.0 || ^18.0 dependencies: - '@babel/runtime': registry.npmmirror.com/@babel/runtime/7.20.13 + '@babel/runtime': 7.20.13 '@radix-ui/react-use-callback-ref': registry.npmmirror.com/@radix-ui/react-use-callback-ref/1.0.0_react@17.0.2 react: registry.npmmirror.com/react/17.0.2 dev: false @@ -4255,7 +4311,7 @@ packages: peerDependencies: react: ^16.8 || ^17.0 || ^18.0 dependencies: - '@babel/runtime': registry.npmmirror.com/@babel/runtime/7.20.13 + '@babel/runtime': 7.20.13 '@radix-ui/rect': registry.npmmirror.com/@radix-ui/rect/1.0.0 react: registry.npmmirror.com/react/17.0.2 dev: false @@ -4280,7 +4336,7 @@ packages: peerDependencies: react: ^16.8 || ^17.0 || ^18.0 dependencies: - '@babel/runtime': registry.npmmirror.com/@babel/runtime/7.20.13 + '@babel/runtime': 7.20.13 '@radix-ui/react-use-layout-effect': registry.npmmirror.com/@radix-ui/react-use-layout-effect/1.0.0_react@17.0.2 react: registry.npmmirror.com/react/17.0.2 dev: false @@ -4318,7 +4374,7 @@ packages: name: '@radix-ui/rect' version: 0.1.1 dependencies: - '@babel/runtime': registry.npmmirror.com/@babel/runtime/7.20.13 + '@babel/runtime': 7.20.13 dev: false registry.npmmirror.com/@radix-ui/rect/1.0.0: @@ -4326,7 +4382,7 @@ packages: name: '@radix-ui/rect' version: 1.0.0 dependencies: - '@babel/runtime': registry.npmmirror.com/@babel/runtime/7.20.13 + '@babel/runtime': 7.20.13 dev: false registry.npmmirror.com/@rollup/plugin-babel/5.3.1_3dsfpkpoyvuuxyfgdbpn4j4uzm: @@ -12867,7 +12923,7 @@ packages: name: regenerator-transform version: 0.15.1 dependencies: - '@babel/runtime': registry.npmmirror.com/@babel/runtime/7.20.13 + '@babel/runtime': 7.20.13 dev: false registry.npmmirror.com/regex-parser/2.2.11: @@ -13088,7 +13144,7 @@ packages: name: rtl-css-js version: 1.16.1 dependencies: - '@babel/runtime': registry.npmmirror.com/@babel/runtime/7.20.13 + '@babel/runtime': 7.20.13 dev: false registry.npmmirror.com/run-parallel/1.2.0: diff --git a/lama_cleaner/app/src/App.tsx b/lama_cleaner/app/src/App.tsx index b771bdb..f4e8cb7 100644 --- a/lama_cleaner/app/src/App.tsx +++ b/lama_cleaner/app/src/App.tsx @@ -9,6 +9,7 @@ import { keepGUIAlive } from './utils' import Header from './components/Header/Header' import useHotKey from './hooks/useHotkey' import { getServerConfig, isDesktop } from './adapters/inpainting' +import SDProgress from './components/SDProgress/SDProgress' const SUPPORTED_FILE_TYPE = [ 'image/jpeg', @@ -178,6 +179,7 @@ function App() { return (
+
) diff --git a/lama_cleaner/app/src/components/SDProgress/SDProgress.scss b/lama_cleaner/app/src/components/SDProgress/SDProgress.scss new file mode 100644 index 0000000..52dde6e --- /dev/null +++ b/lama_cleaner/app/src/components/SDProgress/SDProgress.scss @@ -0,0 +1,35 @@ +.ProgressWrapper { + position: fixed; + left: 50%; + transform: translateX(-50%); + top: 68px; + height: 32px; + display: flex; + justify-content: center; + align-items: center; + gap: 18px; + border: 1px solid var(--border-color); + border-radius: 14px; + padding-left: 8px; + padding-right: 8px; +} + +.ProgressRoot { + position: relative; + overflow: hidden; + background: var(--page-bg); + border-radius: 99999px; + width: 130px; + height: 10px; + + /* Fix overflow clipping in Safari */ + /* https://gist.github.com/domske/b66047671c780a238b51c51ffde8d3a0 */ + transform: translateZ(0); +} + +.ProgressIndicator { + background-color: var(--yellow-accent); + width: 100%; + height: 100%; + transition: transform 60ms cubic-bezier(0.65, 0, 0.35, 1); +} diff --git a/lama_cleaner/app/src/components/SDProgress/SDProgress.tsx b/lama_cleaner/app/src/components/SDProgress/SDProgress.tsx new file mode 100644 index 0000000..f439b6b --- /dev/null +++ b/lama_cleaner/app/src/components/SDProgress/SDProgress.tsx @@ -0,0 +1,77 @@ +import * as React from 'react' +import * as Progress from '@radix-ui/react-progress' +import { useRecoilValue } from 'recoil' +import io from 'socket.io-client' +import { isInpaintingState, isSDState, settingState } from '../../store/Atoms' + +const isDev = process.env.NODE_ENV === 'development' +const socket = isDev ? io(`http://localhost:8080`) : io() + +const SDProgress = () => { + const isSD = useRecoilValue(isSDState) + const isInpainting = useRecoilValue(isInpaintingState) + const [isConnected, setIsConnected] = React.useState(false) + const [step, setStep] = React.useState(0) + const setting = useRecoilValue(settingState) + const maxStep = Math.max(setting.sdSteps, 1) + + const progress = Math.min(Math.round((step / maxStep) * 100), 100) + + React.useEffect(() => { + if (!isSD) return + + socket.on('connect', () => { + setIsConnected(true) + }) + + socket.on('disconnect', () => { + setIsConnected(false) + }) + + socket.on('diffusion_progress', data => { + console.log(`step: ${data.step + 1}`) + if (data) { + setStep(data.step + 1) + } + }) + + socket.on('diffusion_finish', data => { + setStep(0) + }) + + return () => { + socket.off('connect') + socket.off('disconnect') + socket.off('diffusion_progress') + socket.off('diffusion_finish') + } + }, [isSD]) + + return ( +
+ + + +
+ {progress}% +
+
+ ) +} + +export default SDProgress diff --git a/lama_cleaner/app/src/styles/_index.scss b/lama_cleaner/app/src/styles/_index.scss index 6774a6e..c781b50 100644 --- a/lama_cleaner/app/src/styles/_index.scss +++ b/lama_cleaner/app/src/styles/_index.scss @@ -19,6 +19,7 @@ @use '../components/Plugins/Plugins.scss'; @use '../components/Croper/Croper.scss'; @use '../components/InteractiveSeg/InteractiveSeg.scss'; +@use '../components/SDProgress/SDProgress.scss'; // Shared @use '../components/FileSelect/FileSelect'; diff --git a/lama_cleaner/server.py b/lama_cleaner/server.py index a36bb91..f6509fb 100644 --- a/lama_cleaner/server.py +++ b/lama_cleaner/server.py @@ -1,4 +1,5 @@ #!/usr/bin/env python3 +import asyncio import hashlib import os @@ -49,6 +50,7 @@ from flask import ( send_from_directory, jsonify, ) +from flask_socketio import SocketIO # Disable ability for Flask to display warning about using a development server in a production environment. # https://gist.github.com/jerblack/735b9953ba1ab6234abb43174210d356 @@ -88,6 +90,7 @@ logging.getLogger("werkzeug").addFilter(NoFlaskwebgui()) app = Flask(__name__, static_folder=os.path.join(BUILD_DIR, "static")) app.config["JSON_AS_ASCII"] = False CORS(app, expose_headers=["Content-Disposition"]) +socketio = SocketIO(app, cors_allowed_origins="*", async_mode='threading') model: ModelManager = None thumb: FileManager = None @@ -111,8 +114,7 @@ def get_image_ext(img_bytes): def diffuser_callback(i, t, latents): - pass - # socketio.emit('diffusion_step', {'diffusion_step': step}) + socketio.emit('diffusion_progress', {'step': i}) @app.route("/save_image", methods=["POST"]) @@ -301,6 +303,8 @@ def process(): ) ) response.headers["X-Seed"] = str(config.sd_seed) + + socketio.emit('diffusion_finish') return response @@ -557,6 +561,7 @@ def main(args): ui = FlaskUI( app, + socketio=socketio, width=app_width, height=app_height, host=args.host, @@ -565,4 +570,4 @@ def main(args): ) ui.run() else: - app.run(host=args.host, port=args.port, debug=args.debug) + socketio.run(app, host=args.host, port=args.port, debug=args.debug) diff --git a/requirements.txt b/requirements.txt index 718bc44..8e4a89a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ torch>=1.9.0 opencv-python flask==2.2.3 +flask-socketio flask_cors flaskwebgui==0.3.5 pydantic @@ -12,4 +13,4 @@ transformers==4.27.4 gradio piexif==1.1.3 safetensors -omegaconf +omegaconf \ No newline at end of file