Update app_zero.py

app_zero.py

@@ -3,13 +3,23 @@ import os
 import types
 import random
 import datetime
+from pathlib import Path
+
+import huggingface_hub
+
+# -------------------------------------------------------------------
+# Compatibility shim: older diffusers may still expect cached_download
+# -------------------------------------------------------------------
+if not hasattr(huggingface_hub, "cached_download"):
+    def cached_download(*args, **kwargs):
+        return huggingface_hub.hf_hub_download(*args, **kwargs)
+    huggingface_hub.cached_download = cached_download
 
 import torch
 import numpy as np
 import einops
 import spaces
 import gradio as gr
-import huggingface_hub
 
 from PIL import Image
 from torchvision import transforms
@@ -26,7 +36,9 @@ from diffusers import (
     UniPCMultistepScheduler,
 )
 
-# ---- GPU spoof ----
+# -------------------------------------------------------------------
+# GPU spoof for Spaces env compatibility
+# -------------------------------------------------------------------
 torch.cuda.get_device_capability = lambda *args, **kwargs: (8, 6)
 torch.cuda.get_device_properties = lambda *args, **kwargs: types.SimpleNamespace(
     name="NVIDIA A10G",
@@ -36,7 +48,9 @@ torch.cuda.get_device_properties = lambda *args, **kwargs: types.SimpleNamespace
     multi_processor_count=80,
 )
 
-# ---- Downloads ----
+# -------------------------------------------------------------------
+# Download required assets
+# -------------------------------------------------------------------
 huggingface_hub.snapshot_download(
     repo_id="camenduru/PASD",
     allow_patterns=[
@@ -60,48 +74,168 @@ huggingface_hub.hf_hub_download(
     local_dir="PASD/annotator/ckpts",
 )
 
-# ---- PASD ----
+# -------------------------------------------------------------------
+# PASD local path
+# -------------------------------------------------------------------
 sys.path.append("./PASD")
 
+
+# -------------------------------------------------------------------
+# Runtime patching for PASD legacy imports
+# -------------------------------------------------------------------
+def patch_file(path_str: str, replacements: list[tuple[str, str]]) -> None:
+    path = Path(path_str)
+    if not path.exists():
+        print(f"[patch] file not found: {path}")
+        return
+
+    try:
+        text = path.read_text(encoding="utf-8")
+    except Exception as e:
+        print(f"[patch] failed reading {path}: {e}")
+        return
+
+    original = text
+
+    for old, new in replacements:
+        text = text.replace(old, new)
+
+    if text != original:
+        try:
+            path.write_text(text, encoding="utf-8")
+            print(f"[patch] updated: {path}")
+        except Exception as e:
+            print(f"[patch] failed writing {path}: {e}")
+    else:
+        print(f"[patch] no changes: {path}")
+
+
+def patch_pasd_for_diffusers() -> None:
+    # 1) pipeline_utils path moved
+    patch_file(
+        "./PASD/pipelines/pipeline_pasd.py",
+        [
+            (
+                "from diffusers.pipeline_utils import DiffusionPipeline",
+                "from diffusers import DiffusionPipeline",
+            ),
+        ],
+    )
+
+    # 2) PositionNet was renamed/replaced by GLIGENTextBoundingboxProjection
+    #    This patch handles the common legacy multiline import block.
+    patch_file(
+        "./PASD/models/pasd/unet_2d_condition.py",
+        [
+            (
+                "    PositionNet,\n",
+                "",
+            ),
+            (
+                "    GLIGENTextBoundingboxProjection,\n",
+                "    GLIGENTextBoundingboxProjection as PositionNet,\n",
+            ),
+        ],
+    )
+
+    # 3) internal module paths moved in newer diffusers
+    patch_file(
+        "./PASD/models/pasd/unet_2d_blocks.py",
+        [
+            (
+                "from diffusers.models.attention import AdaGroupNorm",
+                "from diffusers.models.normalization import AdaGroupNorm",
+            ),
+            (
+                "from diffusers.models.dual_transformer_2d import DualTransformer2DModel",
+                "from diffusers.models.transformers.dual_transformer_2d import DualTransformer2DModel",
+            ),
+            (
+                "from diffusers.models.transformer_2d import Transformer2DModel",
+                "from diffusers.models.transformers.transformer_2d import Transformer2DModel",
+            ),
+        ],
+    )
+
+    # 4) loader mixin path/name changed across diffusers versions
+    patch_file(
+        "./PASD/models/pasd/controlnet.py",
+        [
+            (
+                "from diffusers.loaders import FromOriginalControlnetMixin",
+                "from diffusers.loaders.single_file_model import FromOriginalModelMixin as FromOriginalControlnetMixin",
+            ),
+        ],
+    )
+
+
+patch_pasd_for_diffusers()
+
+# -------------------------------------------------------------------
+# Import PASD modules only after patching
+# -------------------------------------------------------------------
 from pipelines.pipeline_pasd import StableDiffusionControlNetPipeline
 from myutils.misc import load_dreambooth_lora
 from myutils.wavelet_color_fix import wavelet_color_fix
 from annotator.retinaface import RetinaFaceDetection
 
-from models.pasd.unet_2d_condition import UNet2DConditionModel
-from models.pasd.controlnet import ControlNetModel
+use_pasd_light = False
+face_detector = RetinaFaceDetection()
 
-# ---- Setup ----
+if use_pasd_light:
+    from models.pasd_light.unet_2d_condition import UNet2DConditionModel
+    from models.pasd_light.controlnet import ControlNetModel
+else:
+    from models.pasd.unet_2d_condition import UNet2DConditionModel
+    from models.pasd.controlnet import ControlNetModel
+
+# -------------------------------------------------------------------
+# Model setup
+# -------------------------------------------------------------------
 pretrained_model_path = "stable-diffusion-v1-5/stable-diffusion-v1-5"
 ckpt_path = "PASD/runs/pasd/checkpoint-100000"
 dreambooth_lora_path = "PASD/checkpoints/personalized_models/majicmixRealistic_v6.safetensors"
 
-device = "cuda"
 weight_dtype = torch.float16
+device = "cuda"
 
-# ---- Load ----
-scheduler = UniPCMultistepScheduler.from_pretrained(pretrained_model_path, subfolder="scheduler")
-text_encoder = CLIPTextModel.from_pretrained(pretrained_model_path, subfolder="text_encoder")
-tokenizer = CLIPTokenizer.from_pretrained(pretrained_model_path, subfolder="tokenizer")
-vae = AutoencoderKL.from_pretrained(pretrained_model_path, subfolder="vae")
-feature_extractor = CLIPImageProcessor.from_pretrained(pretrained_model_path, subfolder="feature_extractor")
-
-unet = UNet2DConditionModel.from_pretrained(ckpt_path, subfolder="unet")
-controlnet = ControlNetModel.from_pretrained(ckpt_path, subfolder="controlnet")
+scheduler = UniPCMultistepScheduler.from_pretrained(
+    pretrained_model_path, subfolder="scheduler"
+)
+text_encoder = CLIPTextModel.from_pretrained(
+    pretrained_model_path, subfolder="text_encoder"
+)
+tokenizer = CLIPTokenizer.from_pretrained(
+    pretrained_model_path, subfolder="tokenizer"
+)
+vae = AutoencoderKL.from_pretrained(
+    pretrained_model_path, subfolder="vae"
+)
+feature_extractor = CLIPImageProcessor.from_pretrained(
+    pretrained_model_path, subfolder="feature_extractor"
+)
+unet = UNet2DConditionModel.from_pretrained(
+    ckpt_path, subfolder="unet"
+)
+controlnet = ControlNetModel.from_pretrained(
+    ckpt_path, subfolder="controlnet"
+)
 
 vae.requires_grad_(False)
 text_encoder.requires_grad_(False)
 unet.requires_grad_(False)
 controlnet.requires_grad_(False)
 
-unet, vae, text_encoder = load_dreambooth_lora(unet, vae, text_encoder, dreambooth_lora_path)
+unet, vae, text_encoder = load_dreambooth_lora(
+    unet, vae, text_encoder, dreambooth_lora_path
+)
 
 text_encoder.to(device, dtype=weight_dtype)
 vae.to(device, dtype=weight_dtype)
 unet.to(device, dtype=weight_dtype)
 controlnet.to(device, dtype=weight_dtype)
 
-pipeline = StableDiffusionControlNetPipeline(
+validation_pipeline = StableDiffusionControlNetPipeline(
     vae=vae,
     text_encoder=text_encoder,
     tokenizer=tokenizer,
@@ -113,23 +247,36 @@ pipeline = StableDiffusionControlNetPipeline(
     requires_safety_checker=False,
 )
 
-pipeline._init_tiled_vae(decoder_tile_size=224)
+validation_pipeline._init_tiled_vae(decoder_tile_size=224)
 
-# ---- ResNet ----
+# -------------------------------------------------------------------
+# ResNet helper
+# -------------------------------------------------------------------
 weights = ResNet50_Weights.DEFAULT
 preprocess = weights.transforms()
 resnet = resnet50(weights=weights)
 resnet.eval()
 
-# ---- Utils ----
-def resize_image(image_path, target_height):
+
+def resize_image(image_path: str, target_height: int) -> Image.Image:
     with Image.open(image_path) as img:
         ratio = target_height / float(img.size[1])
         new_width = int(float(img.size[0]) * ratio)
         return img.resize((new_width, target_height), Image.LANCZOS)
 
+
 @spaces.GPU(enable_queue=True)
-def inference(input_image, prompt, a_prompt, n_prompt, steps, upscale, alpha, cfg, seed):
+def inference(
+    input_image,
+    prompt,
+    a_prompt,
+    n_prompt,
+    denoise_steps,
+    upscale,
+    alpha,
+    cfg,
+    seed,
+):
     if seed == -1:
         seed = 0
 
@@ -138,35 +285,53 @@ def inference(input_image, prompt, a_prompt, n_prompt, steps, upscale, alpha, cf
 
     with torch.no_grad():
         seed_everything(seed)
-        generator = torch.Generator(device=device).manual_seed(seed)
+        generator = torch.Generator(device=device)
+        generator.manual_seed(seed)
 
         input_image = input_image.convert("RGB")
 
-        prompt = a_prompt if prompt == "" else f"{prompt}, {a_prompt}"
+        batch = preprocess(input_image).unsqueeze(0)
+        prediction = resnet(batch).squeeze(0).softmax(0)
+        class_id = prediction.argmax().item()
+        score = prediction[class_id].item()
+        category_name = weights.meta["categories"][class_id]
 
-        ori_w, ori_h = input_image.size
-        rscale = upscale
+        if score >= 0.1:
+            prompt += f"{category_name}" if prompt == "" else f", {category_name}"
 
-        input_image = input_image.resize((input_image.size[0]*rscale, input_image.size[1]*rscale))
-        input_image = input_image.resize((input_image.size[0]//8*8, input_image.size[1]//8*8))
+        prompt = a_prompt if prompt == "" else f"{prompt}, {a_prompt}"
 
-        w, h = input_image.size
+        ori_width, ori_height = input_image.size
 
-        image = pipeline(
-            None,
-            prompt,
-            input_image,
-            num_inference_steps=steps,
-            generator=generator,
-            height=h,
-            width=w,
-            guidance_scale=cfg,
-            negative_prompt=n_prompt,
-            conditioning_scale=alpha,
-        ).images[0]
-
-        image = wavelet_color_fix(image, input_image)
-        image = image.resize((ori_w*rscale, ori_h*rscale))
+        rscale = upscale
+        input_image = input_image.resize(
+            (input_image.size[0] * rscale, input_image.size[1] * rscale)
+        )
+        input_image = input_image.resize(
+            (input_image.size[0] // 8 * 8, input_image.size[1] // 8 * 8)
+        )
+        width, height = input_image.size
+
+        try:
+            image = validation_pipeline(
+                None,
+                prompt,
+                input_image,
+                num_inference_steps=denoise_steps,
+                generator=generator,
+                height=height,
+                width=width,
+                guidance_scale=cfg,
+                negative_prompt=n_prompt,
+                conditioning_scale=alpha,
+                eta=0.0,
+            ).images[0]
+
+            image = wavelet_color_fix(image, input_image)
+            image = image.resize((ori_width * rscale, ori_height * rscale))
+        except Exception as e:
+            print(f"[inference] error: {e}")
+            image = Image.new(mode="RGB", size=(512, 512))
 
     result_path = f"result_{timestamp}.jpg"
     input_path = f"input_{timestamp}.jpg"
@@ -176,30 +341,128 @@ def inference(input_image, prompt, a_prompt, n_prompt, steps, upscale, alpha, cf
 
     return input_path, result_path, result_path
 
-# ---- UI ----
+
+css = """
+#col-container{
+    margin: 0 auto;
+    max-width: 720px;
+}
+#project-links{
+    margin: 0 0 12px !important;
+    column-gap: 8px;
+    display: flex;
+    justify-content: center;
+    flex-wrap: nowrap;
+    flex-direction: row;
+    align-items: center;
+}
+"""
+
 with gr.Blocks() as demo:
-    with gr.Row():
-        with gr.Column():
-            input_image = gr.Image(type="filepath")
-            prompt = gr.Textbox(label="Prompt")
-
-            steps = gr.Slider(10, 50, 20)
-            upscale = gr.Slider(1, 4, 2)
-            alpha = gr.Slider(0.5, 1.5, 1.1)
-            cfg = gr.Slider(0.1, 10.0, 7.5)
-            seed = gr.Slider(-1, 2147483647, randomize=True)
-
-            btn = gr.Button("Generate")
-
-        with gr.Column():
-            before = gr.Image()
-            after = gr.Image()
-            file = gr.File()
-
-    btn.click(
-        inference,
-        inputs=[input_image, prompt, prompt, prompt, steps, upscale, alpha, cfg, seed],
-        outputs=[before, after, file],
+    with gr.Column(elem_id="col-container"):
+        gr.HTML("""
+        <h2 style="text-align: center;">
+            PASD Magnify
+        </h2>
+        <p style="text-align: center;">
+            Pixel-Aware Stable Diffusion for Realistic Image Super-resolution and Personalized Stylization
+        </p>
+        <p id="project-links" align="center">
+            <a href="https://github.com/yangxy/PASD"><img src="https://img.shields.io/badge/Project-Page-Green"></a>
+            <a href="https://huggingface.co/papers/2308.14469"><img src="https://img.shields.io/badge/Paper-Arxiv-red"></a>
+        </p>
+        <p style="margin:12px auto;display: flex;justify-content: center;">
+            <a href="https://huggingface.co/spaces/fffiloni/PASD?duplicate=true">
+                <img src="https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-lg.svg" alt="Duplicate this Space">
+            </a>
+        </p>
+        """)
+
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(
+                    type="filepath",
+                    sources=["upload"],
+                    value="PASD/samples/frog.png",
+                    label="Input image",
+                )
+                prompt_in = gr.Textbox(label="Prompt", value="Frog")
+
+                with gr.Accordion(label="Advanced settings", open=False):
+                    added_prompt = gr.Textbox(
+                        label="Added Prompt",
+                        value="clean, high-resolution, 8k, best quality, masterpiece",
+                    )
+                    neg_prompt = gr.Textbox(
+                        label="Negative Prompt",
+                        value="dotted, noise, blur, lowres, oversmooth, longbody, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality",
+                    )
+                    denoise_steps = gr.Slider(
+                        label="Denoise Steps",
+                        minimum=10,
+                        maximum=50,
+                        value=20,
+                        step=1,
+                    )
+                    upsample_scale = gr.Slider(
+                        label="Upsample Scale",
+                        minimum=1,
+                        maximum=4,
+                        value=2,
+                        step=1,
+                    )
+                    condition_scale = gr.Slider(
+                        label="Conditioning Scale",
+                        minimum=0.5,
+                        maximum=1.5,
+                        value=1.1,
+                        step=0.1,
+                    )
+                    classifier_free_guidance = gr.Slider(
+                        label="Classifier-free Guidance",
+                        minimum=0.1,
+                        maximum=10.0,
+                        value=7.5,
+                        step=0.1,
+                    )
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=-1,
+                        maximum=2147483647,
+                        step=1,
+                        randomize=True,
+                    )
+
+                submit_btn = gr.Button("Submit")
+
+            with gr.Column():
+                before_img = gr.Image(label="Input")
+                after_img = gr.Image(label="Result")
+                file_output = gr.File(label="Downloadable image result")
+
+    submit_btn.click(
+        fn=inference,
+        inputs=[
+            input_image,
+            prompt_in,
+            added_prompt,
+            neg_prompt,
+            denoise_steps,
+            upsample_scale,
+            condition_scale,
+            classifier_free_guidance,
+            seed,
+        ],
+        outputs=[
+            before_img,
+            after_img,
+            file_output,
+        ],
+        api_visibility="private",
     )
 
-demo.queue().launch(ssr_mode=False, mcp_server=False)
+demo.queue(max_size=10).launch(
+    ssr_mode=False,
+    mcp_server=False,
+    css=css,
+)