Add-lightweight-UI-GPU-ping

app.py

@@ -107,8 +107,12 @@ PREVIEW_SIZE = 320
 IMAGE_KEY_CACHE_SIZE = 512
 DEFAULT_RETURN_SCALE = 1
 DEFAULT_BATCH_SIZE = 1
-DEFAULT_UI_BATCH_SIZE = 2
+DEFAULT_UI_BATCH_SIZE = 1
 MAX_BATCH_SIZE = 8
+PING_MODEL_ID = 50
+PING_SIZE = 64
+PING_SAMPLER = "euler"
+PING_SCHEDULER = "simple"
 STARTUP_ASSET_IDS = {
     "checkpoints": (16, 50, 35),
     "loras": (1, 3, 4, 30, 47, 12, 7, 49),
@@ -1226,6 +1230,32 @@ def infer_image(
         return image
 
 
+@spaces.GPU(duration=1)
+def infer_ping(model, clip, vae):
+    with torch.inference_mode():
+        latent = state["latent"].generate(
+            width=PING_SIZE,
+            height=PING_SIZE,
+            batch_size=1,
+        )[0]
+        positive = state["encode"].encode(clip=clip, text="")[0]
+        negative = state["encode"].encode(clip=clip, text="")[0]
+        samples = run_sampler(
+            model,
+            secrets.randbits(64),
+            1,
+            1,
+            PING_SAMPLER,
+            PING_SCHEDULER,
+            positive,
+            negative,
+            latent,
+            1,
+        )
+        image = state["decode"].decode(vae=vae, samples=samples)[0]
+        return tensor_image(image).resize((1, 1), Image.Resampling.LANCZOS)
+
+
 def generate_gpu(request, regions, style_image, second_style_image):
     latent = state["latent"].generate(
         width=request.width,
@@ -2341,6 +2371,21 @@ def generate(
     ]
 
 
+def ping_image(p=""):
+    if not valid_pass(p):
+        raise gr.Error("Invalid password")
+
+    with lock:
+        prefix = f"{PING_MODEL_ID}_"
+        models = [name for name in generation_models() if name.startswith(prefix)]
+        if len(models) != 1:
+            raise gr.Error(f"Expected one checkpoint with prefix {prefix}")
+        model_name = models[0]
+        model, clip = load_chain(model_name, [])
+        vae = load_vae(vae_name(model_name))
+        return [infer_ping(model, clip, vae)]
+
+
 def api_generate(body, p=""):
     if not valid_pass(p):
         raise gr.Error("Invalid password")
@@ -3474,7 +3519,9 @@ with gr.Blocks(title="Image generation") as demo:
                     step=.01,
                     label="Return scale",
                 )
-                button = gr.Button("Generate", variant="primary")
+                with gr.Row():
+                    button = gr.Button("Generate", variant="primary")
+                    ping_button = gr.Button("Ping")
             with gr.Column(scale=6, min_width=420):
                 output = gr.Gallery(
                     label="Preview",
@@ -3539,6 +3586,12 @@ with gr.Blocks(title="Image generation") as demo:
         api_name="ui_generate",
         api_visibility="private",
     )
+    ping_button.click(
+        ping_image,
+        pass_input,
+        output,
+        api_visibility="private",
+    )
     upscale_input.change(
         lambda enabled: gr.Column(visible=enabled),
         upscale_input,