Refactor Gradio demo structure and add inference step control

app.py

@@ -1,285 +1,12 @@
-from pathlib import Path
+from gradio_app.patches import apply_runtime_patches
 
-def patch_asyncio_cleanup_error() -> None:
-    try:
-        import asyncio.base_events as base_events
-    except Exception:
-        return
 
-    original_del = getattr(base_events.BaseEventLoop, "__del__", None)
-    if original_del is None:
-        return
+apply_runtime_patches()
 
-    def patched_del(self):
-        try:
-            original_del(self)
-        except ValueError as exc:
-            if "Invalid file descriptor" not in str(exc):
-                raise
+from gradio_app.demo import create_demo
 
-    base_events.BaseEventLoop.__del__ = patched_del
 
-
-patch_asyncio_cleanup_error()
-
-import gradio as gr
-import torch
-from huggingface_hub import hf_hub_download
-from PIL import Image
-
-
-def patch_qwen_diffusers_bug() -> None:
-    import importlib.util
-
-    spec = importlib.util.find_spec("diffusers")
-    if spec is None or spec.origin is None:
-        return
-
-    target_file = (
-        Path(spec.origin).resolve().parent
-        / "pipelines"
-        / "qwenimage"
-        / "pipeline_qwenimage_edit_plus.py"
-    )
-    if not target_file.exists():
-        return
-
-    text = target_file.read_text(encoding="utf-8")
-    match = "if prompt_embeds_mask is not None and prompt_embeds_mask.all()"
-    if f"# {match}" in text:
-        return
-
-    lines = text.splitlines()
-    for idx, line in enumerate(lines):
-        if line.strip() == match:
-            if not lines[idx].lstrip().startswith("#"):
-                lines[idx] = f"# {lines[idx]}"
-            if idx + 1 < len(lines) and not lines[idx + 1].lstrip().startswith("#"):
-                lines[idx + 1] = f"# {lines[idx + 1]}"
-            break
-
-    target_file.write_text("\n".join(lines) + "\n", encoding="utf-8")
-
-
-patch_qwen_diffusers_bug()
-
-from diffusers import QwenImageEditPlusPipeline
-
-from pixelsmile.linear_conditioning import compute_text_embeddings
-from pixelsmile.utils.image import resize
-
-
-SUPPORTED_EXPRESSIONS = [
-    "angry",
-    "confused",
-    "contempt",
-    "confident",
-    "disgust",
-    "fear",
-    "happy",
-    "sad",
-    "shy",
-    "sleepy",
-    "surprised",
-    "anxious",
-]
-
-DEFAULT_METHOD = "score_one_all"
-DEFAULT_INF_STEPS = 50
-DEFAULT_RESIZE_MODE = "crop"
-DEFAULT_WIDTH = 512
-DEFAULT_HEIGHT = 512
-DEFAULT_DATA_TYPE = "human"
-DEFAULT_SEED = 42
-DEFAULT_WEIGHT_VERSION = "preview"
-
-ROOT_DIR = Path(__file__).resolve().parent
-WEIGHTS_DIR = ROOT_DIR / "weights"
-BASE_MODEL_REPO = "Qwen/Qwen-Image-Edit-2511"
-PIXELSMILE_DIR = WEIGHTS_DIR / "PixelSmile"
-PIXELSMILE_REPO = "PixelSmile/PixelSmile"
-WEIGHT_FILES = {
-    "preview": "PixelSmile-preview.safetensors",
-    "stable": "PixelSmile-stable.safetensors",
-}
-
-PIPE = None
-PIPE_STATE = {"version": None, "device": None}
-
-
-def get_subject_name(data_type: str) -> str:
-    if data_type == "human":
-        return "person"
-    if data_type == "anime":
-        return "character"
-    raise ValueError(f"Unsupported data_type: {data_type}")
-
-
-def build_edit_condition(subject: str, expression: str, scale: float) -> dict:
-    return {
-        "prompt": f"Edit the {subject} to show a {expression} expression",
-        "prompt_neu": f"Edit the {subject} to show a neutral expression",
-        "category": expression,
-        "scores": {expression: scale},
-    }
-
-
-def resolve_lora_path(weight_version: str) -> Path:
-    if weight_version not in WEIGHT_FILES:
-        raise ValueError(f"Unsupported weight version: {weight_version}")
-    return PIXELSMILE_DIR / WEIGHT_FILES[weight_version]
-
-
-def ensure_lora_path(weight_version: str) -> Path:
-    PIXELSMILE_DIR.mkdir(parents=True, exist_ok=True)
-    lora_path = resolve_lora_path(weight_version)
-    if lora_path.exists():
-        return lora_path
-
-    filename = WEIGHT_FILES[weight_version]
-    downloaded_path = hf_hub_download(
-        repo_id=PIXELSMILE_REPO,
-        filename=filename,
-        local_dir=str(PIXELSMILE_DIR),
-    )
-    return Path(downloaded_path)
-
-
-def get_device() -> torch.device:
-    return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
-
-def load_pipe(weight_version: str) -> QwenImageEditPlusPipeline:
-    global PIPE
-
-    device = get_device()
-    device_key = str(device)
-    if PIPE is not None and PIPE_STATE["version"] == weight_version and PIPE_STATE["device"] == device_key:
-        return PIPE
-
-    lora_path = ensure_lora_path(weight_version)
-    pipe = QwenImageEditPlusPipeline.from_pretrained(
-        BASE_MODEL_REPO,
-        torch_dtype=torch.bfloat16,
-        cache_dir=str(WEIGHTS_DIR),
-    )
-    pipe.load_lora_weights(str(lora_path))
-    pipe.to(device)
-
-    PIPE = pipe
-    PIPE_STATE["version"] = weight_version
-    PIPE_STATE["device"] = device_key
-    return PIPE
-
-
-def prepare_input_image(image: Image.Image) -> Image.Image:
-    if image is None:
-        raise gr.Error("Please upload an input image.")
-    if not isinstance(image, Image.Image):
-        image = Image.fromarray(image)
-    image = image.convert("RGB")
-    return resize(image, (DEFAULT_WIDTH, DEFAULT_HEIGHT), DEFAULT_RESIZE_MODE)
-
-
-def run_edit(
-    image: Image.Image,
-    expression: str,
-    scale: float,
-    data_type: str,
-    seed: int,
-    weight_version: str,
-) -> Image.Image:
-    subject = get_subject_name(data_type)
-    pipe = load_pipe(weight_version)
-    input_image = prepare_input_image(image)
-    edit_condition = build_edit_condition(subject, expression, float(scale))
-
-    prompt_embeds, prompt_embeds_mask = compute_text_embeddings(
-        method=DEFAULT_METHOD,
-        pipeline=pipe,
-        data=edit_condition,
-        image=input_image,
-        max_sequence_length=1024,
-    )
-
-    generator = torch.Generator(device=pipe.device).manual_seed(int(seed))
-    with torch.no_grad():
-        output = pipe(
-            image=input_image,
-            prompt_embeds=prompt_embeds,
-            prompt_embeds_mask=prompt_embeds_mask,
-            num_inference_steps=DEFAULT_INF_STEPS,
-            true_cfg_scale=0,
-            output_type="pil",
-            generator=generator,
-        )
-    return output.images[0]
-
-
-def run_demo(
-    image: Image.Image,
-    expression: str,
-    scale: float,
-    data_type: str,
-    seed: int,
-    weight_version: str,
-):
-    try:
-        result = run_edit(
-            image=image,
-            expression=expression,
-            scale=scale,
-            data_type=data_type,
-            seed=seed,
-            weight_version=weight_version,
-        )
-        return result
-    except Exception as exc:
-        raise gr.Error(str(exc)) from exc
-
-
-with gr.Blocks(title="PixelSmile Demo") as demo:
-    gr.Markdown("# PixelSmile Demo")
-    gr.Markdown(
-        "Fine-grained facial expression editing with Qwen-Image-Edit-2511 and PixelSmile weights."
-    )
-
-    with gr.Row():
-        with gr.Column(scale=1):
-            input_image = gr.Image(type="pil", label="Input Image")
-            expression = gr.Dropdown(
-                choices=SUPPORTED_EXPRESSIONS,
-                value="happy",
-                label="Target Expression",
-            )
-            scale = gr.Slider(
-                minimum=0.0,
-                maximum=1.5,
-                step=0.1,
-                value=0.8,
-                label="Expression Strength",
-            )
-            data_type = gr.Radio(
-                choices=["human", "anime"],
-                value=DEFAULT_DATA_TYPE,
-                label="Data Type",
-            )
-            weight_version = gr.Radio(
-                choices=["preview", "stable"],
-                value=DEFAULT_WEIGHT_VERSION,
-                label="PixelSmile Weight Version",
-            )
-            seed = gr.Number(value=DEFAULT_SEED, precision=0, label="Seed")
-            run_button = gr.Button("Run Inference", variant="primary")
-
-        with gr.Column(scale=1):
-            output_image = gr.Image(type="pil", label="Edited Image")
-
-    run_button.click(
-        fn=run_demo,
-        inputs=[input_image, expression, scale, data_type, seed, weight_version],
-        outputs=output_image,
-    )
+demo = create_demo()
 
 
 if __name__ == "__main__":

gradio_app/__init__.py

@@ -0,0 +1 @@
+# Gradio application package for PixelSmile Space.

gradio_app/config.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+
+
+SUPPORTED_EXPRESSIONS = [
+    "angry",
+    "confused",
+    "contempt",
+    "confident",
+    "disgust",
+    "fear",
+    "happy",
+    "sad",
+    "shy",
+    "sleepy",
+    "surprised",
+    "anxious",
+]
+
+DEFAULT_METHOD = "score_one_all"
+DEFAULT_INF_STEPS = 50
+DEFAULT_RESIZE_MODE = "crop"
+DEFAULT_WIDTH = 512
+DEFAULT_HEIGHT = 512
+DEFAULT_DATA_TYPE = "human"
+DEFAULT_SEED = 42
+DEFAULT_WEIGHT_VERSION = "preview"
+
+ROOT_DIR = Path(__file__).resolve().parent.parent
+WEIGHTS_DIR = ROOT_DIR / "weights"
+BASE_MODEL_REPO = "Qwen/Qwen-Image-Edit-2511"
+PIXELSMILE_DIR = WEIGHTS_DIR / "PixelSmile"
+PIXELSMILE_REPO = "PixelSmile/PixelSmile"
+WEIGHT_FILES = {
+    "preview": "PixelSmile-preview.safetensors",
+    "stable": "PixelSmile-stable.safetensors",
+}

gradio_app/demo.py

@@ -0,0 +1,97 @@
+import gradio as gr
+from PIL import Image
+
+from gradio_app.config import (
+    DEFAULT_DATA_TYPE,
+    DEFAULT_INF_STEPS,
+    DEFAULT_SEED,
+    DEFAULT_WEIGHT_VERSION,
+    SUPPORTED_EXPRESSIONS,
+)
+from gradio_app.edit import run_edit
+from gradio_app.pipeline import PRELOAD_STATE, start_preload
+
+
+def run_demo(
+    image: Image.Image,
+    expression: str,
+    scale: float,
+    data_type: str,
+    seed: int,
+    weight_version: str,
+    num_inference_steps: int,
+):
+    try:
+        if PRELOAD_STATE["loading"]:
+            raise gr.Error(
+                "The model is still loading. Please wait for the startup preload to finish and try again."
+            )
+        if PRELOAD_STATE["error"] is not None:
+            raise gr.Error(f"Model preload failed: {PRELOAD_STATE['error']}")
+
+        return run_edit(
+            image=image,
+            expression=expression,
+            scale=scale,
+            data_type=data_type,
+            seed=seed,
+            weight_version=weight_version,
+            num_inference_steps=num_inference_steps,
+        )
+    except Exception as exc:
+        raise gr.Error(str(exc)) from exc
+
+
+def create_demo() -> gr.Blocks:
+    with gr.Blocks(title="PixelSmile Demo") as demo:
+        gr.Markdown("# PixelSmile Demo")
+        gr.Markdown(
+            "Fine-grained facial expression editing with Qwen-Image-Edit-2511 and PixelSmile weights."
+        )
+
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image = gr.Image(type="pil", label="Input Image", height=420)
+                expression = gr.Dropdown(
+                    choices=SUPPORTED_EXPRESSIONS,
+                    value="happy",
+                    label="Target Expression",
+                )
+                scale = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.5,
+                    step=0.1,
+                    value=0.8,
+                    label="Expression Strength",
+                )
+                data_type = gr.Dropdown(
+                    choices=["human"],
+                    value=DEFAULT_DATA_TYPE,
+                    label="Data Type",
+                )
+                gr.Markdown("<span style='font-size: 12px;'>Anime editing support is coming soon.</span>")
+                weight_version = gr.Dropdown(
+                    choices=["preview"],
+                    value=DEFAULT_WEIGHT_VERSION,
+                    label="PixelSmile Weight Version",
+                )
+                gr.Markdown("<span style='font-size: 12px;'>Stable weights are coming soon.</span>")
+                seed = gr.Number(value=DEFAULT_SEED, precision=0, label="Seed")
+                num_inference_steps = gr.Number(
+                    value=DEFAULT_INF_STEPS,
+                    precision=0,
+                    label="Inference Steps",
+                )
+                run_button = gr.Button("Run Inference", variant="primary")
+
+            with gr.Column(scale=1):
+                output_image = gr.Image(type="pil", label="Edited Image", height=420)
+
+        run_button.click(
+            fn=run_demo,
+            inputs=[input_image, expression, scale, data_type, seed, weight_version, num_inference_steps],
+            outputs=output_image,
+        )
+
+    start_preload()
+    return demo

gradio_app/edit.py

@@ -0,0 +1,76 @@
+from PIL import Image
+import torch
+import gradio as gr
+
+from gradio_app.config import (
+    DEFAULT_HEIGHT,
+    DEFAULT_INF_STEPS,
+    DEFAULT_METHOD,
+    DEFAULT_RESIZE_MODE,
+    DEFAULT_WIDTH,
+)
+from gradio_app.pipeline import load_lora
+from pixelsmile.linear_conditioning import compute_text_embeddings
+from pixelsmile.utils.image import resize
+
+
+def get_subject_name(data_type: str) -> str:
+    if data_type == "human":
+        return "person"
+    if data_type == "anime":
+        return "character"
+    raise ValueError(f"Unsupported data_type: {data_type}")
+
+
+def build_edit_condition(subject: str, expression: str, scale: float) -> dict:
+    return {
+        "prompt": f"Edit the {subject} to show a {expression} expression",
+        "prompt_neu": f"Edit the {subject} to show a neutral expression",
+        "category": expression,
+        "scores": {expression: scale},
+    }
+
+
+def prepare_input_image(image: Image.Image) -> Image.Image:
+    if image is None:
+        raise gr.Error("Please upload an input image.")
+    if not isinstance(image, Image.Image):
+        image = Image.fromarray(image)
+    image = image.convert("RGB")
+    return resize(image, (DEFAULT_WIDTH, DEFAULT_HEIGHT), DEFAULT_RESIZE_MODE)
+
+
+def run_edit(
+    image: Image.Image,
+    expression: str,
+    scale: float,
+    data_type: str,
+    seed: int,
+    weight_version: str,
+    num_inference_steps: int,
+) -> Image.Image:
+    subject = get_subject_name(data_type)
+    pipe = load_lora(weight_version)
+    input_image = prepare_input_image(image)
+    edit_condition = build_edit_condition(subject, expression, float(scale))
+
+    prompt_embeds, prompt_embeds_mask = compute_text_embeddings(
+        method=DEFAULT_METHOD,
+        pipeline=pipe,
+        data=edit_condition,
+        image=input_image,
+        max_sequence_length=1024,
+    )
+
+    generator = torch.Generator(device=pipe.device).manual_seed(int(seed))
+    with torch.no_grad():
+        output = pipe(
+            image=input_image,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            num_inference_steps=int(num_inference_steps),
+            true_cfg_scale=0,
+            output_type="pil",
+            generator=generator,
+        )
+    return output.images[0]

gradio_app/patches.py

@@ -0,0 +1,66 @@
+import os
+from pathlib import Path
+
+
+def patch_asyncio_cleanup_error() -> None:
+    try:
+        import asyncio.base_events as base_events
+    except Exception:
+        return
+
+    original_del = getattr(base_events.BaseEventLoop, "__del__", None)
+    if original_del is None:
+        return
+
+    def patched_del(self):
+        try:
+            original_del(self)
+        except ValueError as exc:
+            if "Invalid file descriptor" not in str(exc):
+                raise
+
+    base_events.BaseEventLoop.__del__ = patched_del
+
+
+def configure_hf_download_env() -> None:
+    os.environ.setdefault("HF_HUB_DOWNLOAD_TIMEOUT", "1800")
+    os.environ.setdefault("HF_HUB_ETAG_TIMEOUT", "1800")
+
+
+def patch_qwen_diffusers_bug() -> None:
+    import importlib.util
+
+    spec = importlib.util.find_spec("diffusers")
+    if spec is None or spec.origin is None:
+        return
+
+    target_file = (
+        Path(spec.origin).resolve().parent
+        / "pipelines"
+        / "qwenimage"
+        / "pipeline_qwenimage_edit_plus.py"
+    )
+    if not target_file.exists():
+        return
+
+    text = target_file.read_text(encoding="utf-8")
+    match = "if prompt_embeds_mask is not None and prompt_embeds_mask.all()"
+    if f"# {match}" in text:
+        return
+
+    lines = text.splitlines()
+    for idx, line in enumerate(lines):
+        if line.strip() == match:
+            if not lines[idx].lstrip().startswith("#"):
+                lines[idx] = f"# {lines[idx]}"
+            if idx + 1 < len(lines) and not lines[idx + 1].lstrip().startswith("#"):
+                lines[idx + 1] = f"# {lines[idx + 1]}"
+            break
+
+    target_file.write_text("\n".join(lines) + "\n", encoding="utf-8")
+
+
+def apply_runtime_patches() -> None:
+    patch_asyncio_cleanup_error()
+    configure_hf_download_env()
+    patch_qwen_diffusers_bug()

gradio_app/pipeline.py

@@ -0,0 +1,113 @@
+import threading
+from pathlib import Path
+
+import torch
+from diffusers import QwenImageEditPlusPipeline
+from huggingface_hub import hf_hub_download, snapshot_download
+
+from gradio_app.config import (
+    BASE_MODEL_REPO,
+    DEFAULT_WEIGHT_VERSION,
+    PIXELSMILE_DIR,
+    PIXELSMILE_REPO,
+    WEIGHTS_DIR,
+    WEIGHT_FILES,
+)
+
+
+PIPE = None
+PIPE_STATE = {"version": None, "device": None}
+PRELOAD_STATE = {"loading": False, "ready": False, "error": None}
+
+
+def resolve_lora_path(weight_version: str) -> Path:
+    if weight_version not in WEIGHT_FILES:
+        raise ValueError(f"Unsupported weight version: {weight_version}")
+    return PIXELSMILE_DIR / WEIGHT_FILES[weight_version]
+
+
+def ensure_lora_path(weight_version: str) -> Path:
+    PIXELSMILE_DIR.mkdir(parents=True, exist_ok=True)
+    lora_path = resolve_lora_path(weight_version)
+    if lora_path.exists():
+        return lora_path
+
+    filename = WEIGHT_FILES[weight_version]
+    downloaded_path = hf_hub_download(
+        repo_id=PIXELSMILE_REPO,
+        filename=filename,
+        local_dir=str(PIXELSMILE_DIR),
+    )
+    return Path(downloaded_path)
+
+
+def get_device() -> torch.device:
+    return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+def load_pipe() -> QwenImageEditPlusPipeline:
+    global PIPE
+
+    device = get_device()
+    device_key = str(device)
+    if PIPE is not None and PIPE_STATE["device"] == device_key:
+        return PIPE
+
+    try:
+        model_path = snapshot_download(
+            repo_id=BASE_MODEL_REPO,
+            cache_dir=str(WEIGHTS_DIR),
+            resume_download=True,
+        )
+    except Exception:
+        model_path = snapshot_download(
+            repo_id=BASE_MODEL_REPO,
+            cache_dir=str(WEIGHTS_DIR),
+            local_files_only=True,
+        )
+
+    pipe = QwenImageEditPlusPipeline.from_pretrained(
+        model_path,
+        torch_dtype=torch.bfloat16,
+        cache_dir=str(WEIGHTS_DIR),
+    )
+    pipe.to(device)
+
+    PIPE = pipe
+    PIPE_STATE["version"] = None
+    PIPE_STATE["device"] = device_key
+    return PIPE
+
+
+def load_lora(weight_version: str) -> QwenImageEditPlusPipeline:
+    pipe = load_pipe()
+    device_key = str(get_device())
+    if PIPE_STATE["version"] == weight_version and PIPE_STATE["device"] == device_key:
+        return pipe
+
+    lora_path = ensure_lora_path(weight_version)
+    try:
+        pipe.unload_lora_weights()
+    except AttributeError:
+        pass
+    pipe.load_lora_weights(str(lora_path))
+    PIPE_STATE["version"] = weight_version
+    return pipe
+
+
+def preload_default_pipe() -> None:
+    try:
+        PRELOAD_STATE["loading"] = True
+        PRELOAD_STATE["ready"] = False
+        PRELOAD_STATE["error"] = None
+        load_lora(DEFAULT_WEIGHT_VERSION)
+        PRELOAD_STATE["ready"] = True
+    except Exception as exc:
+        PRELOAD_STATE["error"] = str(exc)
+        print(f"[WARN] Failed to preload PixelSmile pipeline: {exc}")
+    finally:
+        PRELOAD_STATE["loading"] = False
+
+
+def start_preload() -> None:
+    threading.Thread(target=preload_default_pipe, daemon=True).start()

pixelsmile/__init__.py

@@ -1 +1 @@
-# PixelSmile demo package.
+# Shared PixelSmile demo core package.