Update app_zero.py

app_zero.py

@@ -3,7 +3,6 @@ import os
 import types
 import random
 import datetime
-from pathlib import Path
 
 import torch
 import numpy as np
@@ -27,7 +26,7 @@ from diffusers import (
     UniPCMultistepScheduler,
 )
 
-# ---- GPU spoof for Spaces env compatibility ----
+# ---- GPU spoof ----
 torch.cuda.get_device_capability = lambda *args, **kwargs: (8, 6)
 torch.cuda.get_device_properties = lambda *args, **kwargs: types.SimpleNamespace(
     name="NVIDIA A10G",
@@ -37,7 +36,7 @@ torch.cuda.get_device_properties = lambda *args, **kwargs: types.SimpleNamespace
     multi_processor_count=80,
 )
 
-# ---- Download required assets ----
+# ---- Downloads ----
 huggingface_hub.snapshot_download(
     repo_id="camenduru/PASD",
     allow_patterns=[
@@ -61,76 +60,32 @@ huggingface_hub.hf_hub_download(
     local_dir="PASD/annotator/ckpts",
 )
 
-# ---- Local PASD path ----
+# ---- PASD ----
 sys.path.append("./PASD")
 
-# ---- Compat shim for older custom pipelines expecting diffusers.pipeline_utils ----
-try:
-    import diffusers.pipelines.pipeline_utils as _pipeline_utils
-    sys.modules["diffusers.pipeline_utils"] = _pipeline_utils
-except Exception:
-    pass
-
-# ---- Compat shim for older PASD code expecting PositionNet in diffusers.models.embeddings ----
-try:
-    import diffusers.models.embeddings as _embeddings
-    if not hasattr(_embeddings, "PositionNet") and hasattr(_embeddings, "GLIGENTextBoundingboxProjection"):
-        _embeddings.PositionNet = _embeddings.GLIGENTextBoundingboxProjection
-except Exception:
-    pass
-
-# ---- Compat shim for older PASD code expecting AdaGroupNorm in diffusers.models.attention ----
-try:
-    import diffusers.models.attention as _attention
-    from diffusers.models.normalization import AdaGroupNorm as _AdaGroupNorm
-    if not hasattr(_attention, "AdaGroupNorm"):
-        _attention.AdaGroupNorm = _AdaGroupNorm
-except Exception:
-    pass
-
-# ---- Compat shim for older PASD code expecting diffusers.models.dual_transformer_2d ----
-try:
-    import diffusers.models.transformers.dual_transformer_2d as _dual_transformer_2d
-    sys.modules["diffusers.models.dual_transformer_2d"] = _dual_transformer_2d
-except Exception:
-    pass
-
-# ---- Compat shim for older PASD code expecting diffusers.models.transformer_2d ----
-try:
-    import diffusers.models.transformers.transformer_2d as _transformer_2d
-    sys.modules["diffusers.models.transformer_2d"] = _transformer_2d
-except Exception:
-    pass
-
 from pipelines.pipeline_pasd import StableDiffusionControlNetPipeline
-from myutils.misc import load_dreambooth_lora, rand_name
+from myutils.misc import load_dreambooth_lora
 from myutils.wavelet_color_fix import wavelet_color_fix
 from annotator.retinaface import RetinaFaceDetection
 
-# ---- Model selection ----
-use_pasd_light = False
-face_detector = RetinaFaceDetection()
-
-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
+from models.pasd.unet_2d_condition import UNet2DConditionModel
+from models.pasd.controlnet import ControlNetModel
 
+# ---- 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"
 
-weight_dtype = torch.float16
 device = "cuda"
+weight_dtype = torch.float16
 
-# ---- Load models ----
+# ---- 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")
 
@@ -146,7 +101,7 @@ vae.to(device, dtype=weight_dtype)
 unet.to(device, dtype=weight_dtype)
 controlnet.to(device, dtype=weight_dtype)
 
-validation_pipeline = StableDiffusionControlNetPipeline(
+pipeline = StableDiffusionControlNetPipeline(
     vae=vae,
     text_encoder=text_encoder,
     tokenizer=tokenizer,
@@ -158,25 +113,23 @@ validation_pipeline = StableDiffusionControlNetPipeline(
     requires_safety_checker=False,
 )
 
-validation_pipeline._init_tiled_vae(decoder_tile_size=224)
+pipeline._init_tiled_vae(decoder_tile_size=224)
 
-# ---- ResNet auto-tag helper ----
+# ---- ResNet ----
 weights = ResNet50_Weights.DEFAULT
 preprocess = weights.transforms()
 resnet = resnet50(weights=weights)
 resnet.eval()
 
-
+# ---- Utils ----
 def resize_image(image_path, target_height):
     with Image.open(image_path) as img:
         ratio = target_height / float(img.size[1])
         new_width = int(float(img.size[0]) * ratio)
-        resized_img = img.resize((new_width, target_height), Image.LANCZOS)
-        return resized_img
-
+        return img.resize((new_width, target_height), Image.LANCZOS)
 
 @spaces.GPU(enable_queue=True)
-def inference(input_image, prompt, a_prompt, n_prompt, denoise_steps, upscale, alpha, cfg, seed):
+def inference(input_image, prompt, a_prompt, n_prompt, steps, upscale, alpha, cfg, seed):
     if seed == -1:
         seed = 0
 
@@ -185,52 +138,35 @@ def inference(input_image, prompt, a_prompt, n_prompt, denoise_steps, upscale, a
 
     with torch.no_grad():
         seed_everything(seed)
-        generator = torch.Generator(device=device)
-        generator.manual_seed(seed)
+        generator = torch.Generator(device=device).manual_seed(seed)
 
         input_image = input_image.convert("RGB")
-        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]
-
-        if score >= 0.1:
-            prompt += f"{category_name}" if prompt == "" else f", {category_name}"
 
         prompt = a_prompt if prompt == "" else f"{prompt}, {a_prompt}"
 
-        ori_width, ori_height = input_image.size
-        resize_flag = False
-
+        ori_w, ori_h = input_image.size
         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
-        resize_flag = True
-
-        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)
-
-            if resize_flag:
-                image = image.resize((ori_width * rscale, ori_height * rscale))
-        except Exception as e:
-            print(e)
-            image = Image.new(mode="RGB", size=(512, 512))
+
+        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))
+
+        w, h = 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))
 
     result_path = f"result_{timestamp}.jpg"
     input_path = f"input_{timestamp}.jpg"
@@ -240,110 +176,30 @@ def inference(input_image, prompt, a_prompt, n_prompt, denoise_steps, upscale, a
 
     return input_path, result_path, result_path
 
-
-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;
-}
-"""
-
+# ---- UI ----
 with gr.Blocks() as demo:
-    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",
+    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],
     )
 
-demo.queue(max_size=10).launch(
-    ssr_mode=False,
-    mcp_server=False,
-    css=css,
-)
+demo.queue().launch(ssr_mode=False, mcp_server=False)