[Admin maintenance] Migrate grant to ZeroGPU

app.py

@@ -5,6 +5,7 @@ subprocess.run(
         "pip install ./gradio_magicquill-0.0.1-py3-none-any.whl"
     )
 )
+import spaces
 import gradio as gr
 from gradio_magicquill import MagicQuill
 import random
@@ -13,8 +14,7 @@ import numpy as np
 from PIL import Image, ImageOps
 import base64
 import io
-from fastapi import FastAPI, Request
-import uvicorn
+from fastapi import Request
 from MagicQuill import folder_paths
 from MagicQuill.scribble_color_edit import ScribbleColorEditModel
 from gradio_client import Client, handle_file
@@ -28,7 +28,13 @@ import gc
 snapshot_download(repo_id="LiuZichen/MagicQuill-models", repo_type="model", local_dir="models")
 # HF_TOKEN = os.environ.get("HF_TOKEN")
 # The client has been made public. Welcome to duplicate our repo.
-client = Client("LiuZichen/MagicQuillHelper")
+_HELPER_SPACE = "LiuZichen/MagicQuillHelper"
+_client = None
+def client():
+    global _client
+    if _client is None:
+        _client = Client(_HELPER_SPACE)
+    return _client
 scribbleColorEditModel = ScribbleColorEditModel()
 
 def tensor_to_numpy(tensor):
@@ -130,7 +136,7 @@ def guess_prompt_handler(original_image, add_color_image, add_edge_image):
     add_color_image_file.close()
     add_edge_mask_file.close()
     
-    res = client.predict(
+    res = client().predict(
         handle_file(original_image_file.name),
         handle_file(add_color_image_file.name),
         handle_file(add_edge_mask_file.name),
@@ -186,7 +192,20 @@ def generate(ckpt_name, total_mask, original_image, add_color_image, add_edge_im
     
     return final_image_base64
 
+def _generate_duration(x, ckpt_name, negative_prompt, fine_edge, grow_size, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler):
+    try:
+        s = int(steps or 20)
+    except Exception:
+        s = 20
+    return max(60, 3 * s)
+
+
+from MagicQuill.comfy import model_management as _mm
+
+
+@spaces.GPU(duration=_generate_duration)
 def generate_image_handler(x, ckpt_name, negative_prompt, fine_edge, grow_size, edge_strength, color_strength, inpaint_strength, seed, steps, cfg, sampler_name, scheduler):
+    _mm.unload_all_models()
     if seed == -1:
         seed = random.randint(0, 2**32 - 1)
     ms_data = x['from_frontend']
@@ -321,23 +340,19 @@ with gr.Blocks(css=css, head=head) as demo:
         """)
     demo.queue(max_size=20, status_update_rate=0.1)
 
-app = FastAPI()
 
-@app.post("/magic_quill/guess_prompt")
-async def guess_prompt(request: Request):
-    data = await request.json()
-    res = guess_prompt_handler(data['original_image'], data['add_color_image'], data['add_edge_image'])
-    return res
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860, prevent_thread_lock=True)
 
-@app.post("/magic_quill/process_background_img")
-async def process_background_img(request: Request):
-    img = await request.json()
-    resized_img_tensor = load_and_resize_image(img)
-    resized_img_base64 = "data:image/png;base64," + tensor_to_base64(resized_img_tensor)
-    return resized_img_base64
+    @demo.app.post("/magic_quill/guess_prompt")
+    async def guess_prompt(request: Request):
+        data = await request.json()
+        return guess_prompt_handler(data['original_image'], data['add_color_image'], data['add_edge_image'])
 
-app = gr.mount_gradio_app(app, demo, "/")
+    @demo.app.post("/magic_quill/process_background_img")
+    async def process_background_img(request: Request):
+        img = await request.json()
+        resized_img_tensor = load_and_resize_image(img)
+        return "data:image/png;base64," + tensor_to_base64(resized_img_tensor)
 
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
-    # demo.launch()
+    demo.block_thread()

requirements.txt

@@ -14,7 +14,7 @@ anyio==4.4.0
 async-timeout==4.0.3
 attrs==23.2.0
 beautifulsoup4==4.12.3
-bitsandbytes==0.43.3
+bitsandbytes
 certifi==2024.7.4
 cffi==1.16.0
 chardet==5.2.0
@@ -33,7 +33,7 @@ einops-exts==0.0.4
 embreex==2.17.7.post5
 eval-type-backport==0.2.0
 exceptiongroup==1.2.2
-fastapi
+fastapi<0.112
 ffmpy==0.4.0
 filelock==3.15.4
 flatbuffers==24.3.25
@@ -132,11 +132,11 @@ sounddevice==0.4.7
 soupsieve==2.5
 spandrel==0.3.4
 stanza==1.1.1
-starlette
+starlette<0.38
 svg-path==6.3
 svglib==1.5.1
 svgwrite==1.4.3
-sympy==1.13.1
+sympy==1.13.3
 tabulate==0.9.0
 termcolor==2.4.0
 threadpoolctl==3.5.0
@@ -151,7 +151,8 @@ tqdm==4.66.5
 trampoline==0.1.2
 transformers==4.37.2
 trimesh==4.4.3
-triton==2.1.0
+torch==2.8.0
+torchvision==0.23.0
 torchsde==0.2.6
 typer==0.12.5
 typing-extensions==4.12.2