WIP

app.py

@@ -4,6 +4,7 @@ import spaces
 import os
 import random
 
+import subprocess
 import torch
 from PIL import Image
 import cv2
@@ -32,6 +33,8 @@ Before running, set the `HUGGINGFACE_TOKEN` environment variable **or** call
 `login("<YOUR_HF_TOKEN>")` explicitly.
 """
 
+subprocess.run("rm -rf /data-nvme/zerogpu-offload/*", env={}, shell=True)
+
 # --------------------------------------------------
 # Model & pipeline setup
 # --------------------------------------------------
@@ -76,16 +79,23 @@ MAX_SEED = 100
 # --------------------------------------------------
 
 
-def _preview_canny(pil_img: Image.Image) -> Image.Image:
+def _preview_canny(
+    pil_img: Image.Image, canny_threshold_1: int, canny_threshold_2: int
+) -> Image.Image:
     arr = np.array(pil_img.convert("RGB"))
-    edges = cv2.Canny(arr, 100, 200)
+    edges = cv2.Canny(arr, threshold1=canny_threshold_1, threshold2=canny_threshold_2)
     edges_rgb = cv2.cvtColor(edges, cv2.COLOR_GRAY2RGB)
     return Image.fromarray(edges_rgb)
 
 
-def _make_preview(control_image: Image.Image, mode: str) -> Image.Image:
+def _make_preview(
+    control_image: Image.Image,
+    mode: str,
+    canny_threshold_1: int,
+    canny_threshold_2: int,
+) -> Image.Image:
     if mode == "canny":
-        return _preview_canny(control_image)
+        return _preview_canny(control_image, canny_threshold_1, canny_threshold_2)
     # For other modes you can plug in your own visualiser later
     return control_image
 
@@ -105,6 +115,8 @@ def infer(
     randomize_seed: bool,
     guidance_scale: float,
     num_inference_steps: int,
+    canny_threshold_1: int,
+    canny_threshold_2: int,
 ):
     if control_image is None:
         raise gr.Error("Please upload a control image first.")
@@ -115,8 +127,12 @@ def infer(
     gen = torch.Generator(device).manual_seed(seed)
     w, h = control_image.size
 
+    preprocessed = _make_preview(
+        control_image, mode, canny_threshold_1, canny_threshold_2
+    )
+
     result = pipe(
-        prompt=prompt,
+        prompt=preprocessed,
         control_image=[control_image],
         control_mode=[MODE_MAPPING[mode]],
         width=w,
@@ -127,8 +143,7 @@ def infer(
         generator=gen,
     ).images[0]
 
-    preview = _make_preview(control_image, mode)
-    return result, seed, preview
+    return result, seed, preprocessed
 
 
 # --------------------------------------------------
@@ -148,23 +163,23 @@ with gr.Blocks(css=css, elem_id="wrapper") as demo:
         control_image = gr.Image(
             label="Upload a processed control image",
             type="pil",
-            height=512,
+            height=512 + 256,
         )
-        result_image = gr.Image(label="Result", height=512)
-        preview_image = gr.Image(label="Pre‑processed Cond", height=512)
+        result_image = gr.Image(label="Result", height=512 + 256)
+        preview_image = gr.Image(label="Pre‑processed Cond", height=512 + 256)
 
     # ------------ Prompt ------------
-    prompt_txt = gr.Textbox(label="Prompt", value="best quality", lines=1)
+    prompt_txt = gr.Textbox(label="Prompt", value="A beautiful image", lines=1)
 
     # ------------ ControlNet settings ------------
     with gr.Row():
         with gr.Column():
             gr.Markdown("### ControlNet")
             mode_radio = gr.Radio(
-                choices=list(MODE_MAPPING.keys()), value="gray", label="Mode"
+                choices=list(MODE_MAPPING.keys()), value="canny", label="Mode"
             )
             strength_slider = gr.Slider(
-                0.0, 1.0, value=0.5, step=0.01, label="control strength"
+                0.0, 1.0, value=0.8, step=0.01, label="control strength"
             )
         with gr.Column():
             seed_slider = gr.Slider(0, MAX_SEED, step=1, value=42, label="Seed")
@@ -174,6 +189,15 @@ with gr.Blocks(css=css, elem_id="wrapper") as demo:
             )
             steps_slider = gr.Slider(1, 50, step=1, value=24, label="Inference steps")
 
+    with gr.Row():
+        with gr.Column():
+            gr.Markdown("### Preprocess")
+            canny_threshold_1 = gr.Slider(
+                0, 500, step=1, value=100, label="Canny threshold 1"
+            )
+            canny_threshold_2 = gr.Slider(
+                0, 500, step=1, value=200, label="Canny threshold 2"
+            )
     submit_btn = gr.Button("Submit")
 
     submit_btn.click(
@@ -187,6 +211,8 @@ with gr.Blocks(css=css, elem_id="wrapper") as demo:
             randomize_chk,
             guidance_slider,
             steps_slider,
+            canny_threshold_1,
+            canny_threshold_2,
         ],
         outputs=[result_image, seed_slider, preview_image],
     )