added diffqrcoder_wrapper

app.py

@@ -1,7 +1,120 @@
+# app.py
 import gradio as gr
+import spaces
+from diffqrcoder_wrapper import generate_qr_art
 
-def greet(name):
-    return "Hello " + name + "!!"
+DEFAULT_PROMPT = (
+    "whimsical biomimetic blueprint, iridescent inks swirling through "
+    "mechanical petals, soft gears woven with luminescent filigree"
+)
+
+DEFAULT_NEG = "easynegative"
+
+
+@spaces.GPU  # ZeroGPU: attach GPU only for this call
+def infer(
+    url_or_text: str,
+    prompt: str,
+    num_inference_steps: int,
+    controlnet_scale: float,
+    scanning_robust_guidance_scale: float,
+    perceptual_guidance_scale: float,
+    srmpgd_iters: int,
+):
+    # Map 0 → None for SR-MPGD iteration (as in original script)
+    srmpgd_num_iteration = None if srmpgd_iters == 0 else srmpgd_iters
+
+    img = generate_qr_art(
+        url_or_text=url_or_text,
+        prompt=prompt,
+        neg_prompt=DEFAULT_NEG,
+        num_inference_steps=num_inference_steps,
+        qrcode_module_size=20,
+        qrcode_padding=78,
+        controlnet_conditioning_scale=controlnet_scale,
+        scanning_robust_guidance_scale=scanning_robust_guidance_scale,
+        perceptual_guidance_scale=perceptual_guidance_scale,
+        srmpgd_num_iteration=srmpgd_num_iteration,
+        srmpgd_lr=0.1,
+        seed=1,
+    )
+    return img
+
+
+with gr.Blocks() as demo:
+    gr.Markdown(
+        r"""
+# DiffQRCoder – ZeroGPU demo
+
+Generate aesthetic, scanning-robust QR codes using the **DiffQRCoder** pipeline
+([WACV 2025](https://openaccess.thecvf.com/content/WACV2025/html/Liao_DiffQRCoder_Diffusion-Based_Aesthetic_QR_Code_Generation_with_Scanning_Robustness_Guided_WACV_2025_paper.html)) 🚀
+        """
+    )
+
+    with gr.Row():
+        url = gr.Textbox(
+            label="QR contents (URL or text)",
+            value="https://example.com",
+        )
+
+    prompt = gr.Textbox(
+        label="Style prompt",
+        value=DEFAULT_PROMPT,
+        lines=3,
+    )
+
+    with gr.Accordion("Advanced parameters", open=False):
+        steps = gr.Slider(
+            minimum=10,
+            maximum=60,
+            value=40,
+            step=1,
+            label="Diffusion steps (num_inference_steps)",
+        )
+        control_scale = gr.Slider(
+            minimum=0.5,
+            maximum=2.0,
+            value=1.35,
+            step=0.05,
+            label="ControlNet conditioning scale",
+        )
+        srg_scale = gr.Slider(
+            minimum=0,
+            maximum=800,
+            value=500,
+            step=10,
+            label="Scanning-robust guidance scale (srg)",
+        )
+        pg_scale = gr.Slider(
+            minimum=0,
+            maximum=10,
+            value=2,
+            step=0.5,
+            label="Perceptual guidance scale (pg)",
+        )
+        srmpgd_iters = gr.Slider(
+            minimum=0,
+            maximum=64,
+            value=0,
+            step=1,
+            label="SR-MPGD iterations (0 = disabled)",
+        )
+
+    btn = gr.Button("Generate QR Art ✨", variant="primary")
+    out = gr.Image(label="Output QR art", type="pil")
+
+    btn.click(
+        fn=infer,
+        inputs=[
+            url,
+            prompt,
+            steps,
+            control_scale,
+            srg_scale,
+            pg_scale,
+            srmpgd_iters,
+        ],
+        outputs=[out],
+    )
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
 demo.launch()

diffqrcoder/__init__.py

@@ -0,0 +1 @@
+from .pipeline_diffqrcoder import DiffQRCoderPipeline

diffqrcoder/image_processor.py

@@ -0,0 +1,44 @@
+from typing import Optional
+
+import torch
+
+
+IMAGE_MAX_VAL = 255
+
+
+def min_max_normalize(x: torch.Tensor) -> torch.Tensor:
+    return (x - x.min()) / (x.max() - x.min())
+
+
+def convert_to_gray(
+    images: torch.Tensor,
+    cr: float = 0.2999,
+    cg: float = 0.587,
+    cb: float = 0.1114,
+) -> torch.Tensor:
+
+    assert images.shape[1] == 3, \
+        f"The channel of color images must be 3 but get {images.shape[1]}. They are not color images."
+
+    gray_image = cr * images[:, 0] + cg * images[:, 1] + cb * images[:, 2]
+    return gray_image.unsqueeze(1)
+
+
+def image_binarize(
+    image: torch.Tensor,
+    binary_threshold: Optional[float] = None,
+) -> torch.Tensor:
+
+    if image.shape[1] == 3:
+        image = convert_to_gray(image)
+
+    if binary_threshold is None:
+        if image.max() <= 1:
+            binary_threshold = 0.5
+        else:
+            binary_threshold = 0.5 * IMAGE_MAX_VAL
+    return (image > binary_threshold).to(image.dtype)
+
+
+def crop_padding(x: torch.Tensor, padding: int):
+    return x[:, :, padding:-padding, padding:-padding]

diffqrcoder/losses/__init__.py

@@ -0,0 +1,3 @@
+from .perceptual_loss import PerceptualLoss
+from .scanning_robust_loss import ScanningRobustLoss
+from .personalized_code_loss import PersonalizedCodeLoss

diffqrcoder/losses/perceptual_loss.py

@@ -0,0 +1,70 @@
+from collections import namedtuple
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torchvision.models import vgg16
+from torchvision.transforms import Normalize
+
+
+class VGGFeatureExtractor(nn.Module):
+    def __init__(
+        self,
+        requires_grad: bool = False,
+        pretrained_weights: str = "DEFAULT",
+    ) -> None:
+
+        super().__init__()
+        self.norm = Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225],
+        )
+        self.slice_indices = [(0, 4), (4, 9), (9, 16), (16, 23)]
+        self.slices = nn.ModuleList([nn.Sequential() for _ in range(len(self.slice_indices))])
+        self._initialize_slices(pretrained_weights)
+        self.features = namedtuple("Outputs", [f"layer{i}" for i in range(len(self.slice_indices))])
+
+        if not requires_grad:
+            for param in self.parameters():
+                param.requires_grad = False
+
+    def lp_norm(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.normalize(x, p=2.0, dim=1)
+
+    def _initialize_slices(self, pretrained_weights: str = "DEFAULT") -> None:
+        features = vgg16(weights=pretrained_weights).features
+        for slice_idx, (start, end) in enumerate(self.slice_indices):
+            for i in range(start, end):
+                self.slices[slice_idx].add_module(str(i), features[i])
+
+    def forward(self, x: torch.Tensor) -> namedtuple:
+        outputs = []
+        x = self.norm(x)
+        for slice_model in self.slices:
+            x = self.lp_norm(slice_model(x))
+            outputs.append(x)
+        return self.features(*outputs)
+
+
+class PerceptualLoss(nn.Module):
+    def __init__(
+        self,
+        requires_grad: bool = False,
+        pretrained_weights: str = "DEFAULT",
+    ):
+        super(PerceptualLoss, self).__init__()
+        self.extractor = VGGFeatureExtractor(
+            pretrained_weights=pretrained_weights,
+            requires_grad=requires_grad,
+        )
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        return torch.mean(
+            torch.tensor(
+                [
+                    torch.nn.functional.mse_loss(fx, fy)
+                    for fx, fy in zip(self.extractor(x), self.extractor(y))
+                ]
+            )
+        )

diffqrcoder/losses/personalized_code_loss.py

@@ -0,0 +1,42 @@
+import torch
+from torch import nn
+
+from .scanning_robust_loss import ScanningRobustLoss
+
+
+class PersonalizedCodeLoss(nn.Module):
+    def __init__(
+        self,
+        qrcode_image: torch.Tensor,
+        content_image: torch.Tensor,
+        module_size: int = 16,
+        b_thres: float = 50,
+        w_thres: float = 200,
+        b_soft_value: float = 40 / 255,
+        w_soft_value: float = 220 / 255,
+        code_weight: float = 1e12,
+        content_weight: float = 1e8,
+        device: torch.device = "cuda" if torch.cuda.is_available() else "cpu",
+    ):
+        super(PersonalizedCodeLoss, self).__init__()
+        self.code_loss = ScanningRobustLoss(
+            module_size=module_size,
+        ).to(device)
+
+        self.content_image = content_image
+        self.code_weight = code_weight
+        self.content_weight = content_weight
+        self.qrcode_image = qrcode_image
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        code_loss = self.code_loss(x, self.qrcode_image)
+        perceptual_loss = nn.MSELoss()(x, self.content_image)
+        total_loss = (
+            self.code_weight * code_loss + \
+            self.content_weight * perceptual_loss
+        )
+        return {
+            "code": code_loss,
+            "perceptual": perceptual_loss,
+            "total": total_loss
+        }

diffqrcoder/losses/scanning_robust_loss.py

@@ -0,0 +1,145 @@
+import math
+
+import cv2
+import numpy as np
+import torch
+from torch import nn
+
+from ..image_processor import convert_to_gray, image_binarize, min_max_normalize
+
+
+class GaussianFilter(nn.Module):
+    def __init__(self, module_size: int, filter_thres: float = 0.1) -> None:
+        super().__init__()
+        self.module_size = module_size
+        self.filter_thres = filter_thres
+        self.conv = nn.Conv2d(
+            in_channels=1,
+            out_channels=1,
+            kernel_size=module_size,
+            stride=module_size,
+            padding=0,
+            bias=False,
+            groups=1,
+        )
+        self._setup_filter_weights()
+
+    def _setup_filter_weights(self) -> None:
+        filter_1d = cv2.getGaussianKernel(
+            ksize=self.module_size,
+            sigma=1.5,
+            ktype=cv2.CV_32F
+        )
+        filter_2d = filter_1d * filter_1d.T
+        filter_2d = min_max_normalize(filter_2d)
+        filter_2d[filter_2d < self.filter_thres] = .0
+        gaussian_filter_init = torch.tensor(filter_2d, dtype=torch.float32)
+        self.conv.weight = nn.Parameter(
+            gaussian_filter_init.reshape(1, 1, *gaussian_filter_init.shape),
+            requires_grad=False,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+
+class RegionMeanFilter(nn.Module):
+    def __init__(self, module_size: int) -> None:
+        super().__init__()
+        self.module_size = module_size
+        self.conv = nn.Conv2d(
+            in_channels=1,
+            out_channels=1,
+            kernel_size=module_size,
+            stride=module_size,
+            padding=0,
+            bias=None,
+            groups=1,
+        )
+        self._setup_kernel_weights()
+
+    def _setup_kernel_weights(self) -> None:
+        module_center = int(self.module_size / 2)
+        radius = math.ceil(self.module_size / 6)
+        center_filter = torch.zeros((1, 1, self.module_size, self.module_size))
+        center_filter[
+            :, :,
+            module_center-radius : module_center+radius,
+            module_center-radius : module_center+radius,
+        ] = 1.0
+
+        self.conv.weight = nn.Parameter(
+            center_filter / center_filter.sum(),
+            requires_grad=False,
+        )
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+
+class CenterPixelExtractor(nn.Module):
+    def __init__(self, module_size: int) -> None:
+        super().__init__()
+        self.module_size = module_size
+        self.conv = nn.Conv2d(
+            in_channels=1,
+            out_channels=1,
+            kernel_size=module_size,
+            stride=module_size,
+            padding=0,
+            bias=None,
+            groups=1,
+        )
+        self._setup_kernel_weights()
+
+    def _setup_kernel_weights(self) -> None:
+        module_center = int(self.module_size / 2) + 1
+        center_filter = torch.zeros((1, 1, self.module_size, self.module_size))
+        center_filter[:, :, module_center, module_center] = 1.0
+        self.conv.weight = nn.Parameter(center_filter, requires_grad=False)
+
+    @torch.no_grad()
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+
+class QRCodeErrorExtractor(nn.Module):
+    def __init__(self, module_size: int) -> None:
+        super().__init__()
+        self.module_size = module_size
+        self.region_mean_filter = RegionMeanFilter(module_size)
+        self.center_pixel_extractor = CenterPixelExtractor(module_size=module_size)
+
+    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+        x_center_mean = self.region_mean_filter(x)
+        y_center_pixel = self.center_pixel_extractor(y)
+        error_mask = (y_center_pixel == 0) & (x_center_mean > 0.45) | \
+                     (y_center_pixel == 1) & (x_center_mean < 0.65)
+        return error_mask.float()
+
+
+class ScanningRobustLoss(nn.Module):
+    def __init__(self, module_size: int) -> None:
+        super().__init__()
+        self.gaussian_filter = GaussianFilter(module_size=module_size)
+        self.center_filter = RegionMeanFilter(module_size=module_size)
+        self.module_error_extractor = QRCodeErrorExtractor(module_size=module_size)
+
+    def _compute_error(self, image: torch.Tensor, qrcode: torch.Tensor) -> torch.Tensor:
+        gray_image = convert_to_gray(image)
+        error0 = 2 * torch.relu(gray_image - 0.45) * (1 - qrcode)
+        error1 = 2 * torch.relu(0.65 - gray_image) * qrcode
+        return error0 + error1
+
+    def _compute_ealy_stopping_mask(self, image: torch.Tensor, qrcode: torch.Tensor) -> torch.Tensor:
+        return self.module_error_extractor(
+            convert_to_gray(image.clone().detach()),
+            image_binarize(qrcode),
+        )
+
+    def forward(self, image: torch.Tensor, qrcode: torch.Tensor) -> torch.Tensor:
+        error = self._compute_error(image, qrcode)
+        sample_error = self.gaussian_filter(error)
+        mask = self._compute_ealy_stopping_mask(image, qrcode)
+        return torch.mean(sample_error * mask)

diffqrcoder/pipeline_diffqrcoder.py

@@ -0,0 +1,607 @@
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import torch
+from diffusers import StableDiffusionControlNetPipeline
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.image_processor import PipelineImageInput
+from diffusers.models import ControlNetModel
+from diffusers.pipelines.stable_diffusion import StableDiffusionPipelineOutput
+from diffusers.pipelines.controlnet import MultiControlNetModel
+from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
+from diffusers.utils import is_torch_xla_available, deprecate
+from diffusers.utils.torch_utils import is_compiled_module, is_torch_version
+from tqdm import trange
+
+from .image_processor import image_binarize, crop_padding
+from .srpg import ScanningRobustPerceptualGuidance
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+class DiffQRCoderPipeline(StableDiffusionControlNetPipeline):
+    def _run_stage1(
+        self,
+        prompt: Union[str, List[str]] = None,
+        qrcode: PipelineImageInput = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        timesteps: List[int] = None,
+        sigmas: List[float] = None,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+        guess_mode: bool = False,
+        control_guidance_start: Union[float, List[float]] = 0.0,
+        control_guidance_end: Union[float, List[float]] = 1.0,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[
+            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+        ] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        **kwargs,
+    ):
+        return super().__call__(
+            prompt=prompt,
+            image=qrcode,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            timesteps=timesteps,
+            sigmas=sigmas,
+            guidance_scale=guidance_scale,
+            negative_prompt=negative_prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            eta=eta,
+            generator=generator,
+            latents=latents,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            ip_adapter_image=ip_adapter_image,
+            ip_adapter_image_embeds=ip_adapter_image_embeds,
+            output_type=output_type,
+            return_dict=True,
+            cross_attention_kwargs=cross_attention_kwargs,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            guess_mode=guess_mode,
+            control_guidance_start=control_guidance_start,
+            control_guidance_end=control_guidance_end,
+            clip_skip=clip_skip,
+            callback_on_step_end=callback_on_step_end,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            **kwargs,
+        )
+
+    def _run_stage2(
+        self,
+        prompt: Union[str, List[str]] = None,
+        qrcode: PipelineImageInput = None,
+        qrcode_module_size: int = 20,
+        qrcode_padding: int = 78,
+        ref_image: PipelineImageInput = None,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        timesteps: List[int] = None,
+        sigmas: List[float] = None,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+        guess_mode: bool = False,
+        control_guidance_start: Union[float, List[float]] = 0.0,
+        control_guidance_end: Union[float, List[float]] = 1.0,
+        scanning_robust_guidance_scale: int = 500,
+        perceptual_guidance_scale: int = 10,
+        srmpgd_num_iteration: Optional[int] = None,
+        srmpgd_lr: Optional[float] = 0.1,
+        clip_skip: Optional[int] = None,
+        callback_on_step_end: Optional[
+            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+        ] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        **kwargs,
+    ):
+        self.srpg = ScanningRobustPerceptualGuidance(
+            module_size=qrcode_module_size,
+            scanning_robust_guidance_scale=scanning_robust_guidance_scale,
+            perceptual_guidance_scale=perceptual_guidance_scale,
+        ).to(self.device).to(self.dtype)
+
+        callback = kwargs.pop("callback", None)
+        callback_steps = kwargs.pop("callback_steps", None)
+
+        if callback is not None:
+            deprecate(
+                "callback",
+                "1.0.0",
+                "Passing `callback` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+        if callback_steps is not None:
+            deprecate(
+                "callback_steps",
+                "1.0.0",
+                "Passing `callback_steps` as an input argument to `__call__` is deprecated, consider using `callback_on_step_end`",
+            )
+
+        if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+            callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+
+        controlnet = self.controlnet._orig_mod if is_compiled_module(self.controlnet) else self.controlnet
+
+        # align format for control guidance
+        if not isinstance(control_guidance_start, list) and isinstance(control_guidance_end, list):
+            control_guidance_start = len(control_guidance_end) * [control_guidance_start]
+        elif not isinstance(control_guidance_end, list) and isinstance(control_guidance_start, list):
+            control_guidance_end = len(control_guidance_start) * [control_guidance_end]
+        elif not isinstance(control_guidance_start, list) and not isinstance(control_guidance_end, list):
+            mult = len(controlnet.nets) if isinstance(controlnet, MultiControlNetModel) else 1
+            control_guidance_start, control_guidance_end = (
+                mult * [control_guidance_start],
+                mult * [control_guidance_end],
+            )
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            qrcode,
+            callback_steps,
+            negative_prompt,
+            prompt_embeds,
+            negative_prompt_embeds,
+            ip_adapter_image,
+            ip_adapter_image_embeds,
+            controlnet_conditioning_scale,
+            control_guidance_start,
+            control_guidance_end,
+            callback_on_step_end_tensor_inputs,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._clip_skip = clip_skip
+        self._cross_attention_kwargs = cross_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+
+        if isinstance(controlnet, MultiControlNetModel) and isinstance(controlnet_conditioning_scale, float):
+            controlnet_conditioning_scale = [controlnet_conditioning_scale] * len(controlnet.nets)
+
+        global_pool_conditions = (
+            controlnet.config.global_pool_conditions
+            if isinstance(controlnet, ControlNetModel)
+            else controlnet.nets[0].config.global_pool_conditions
+        )
+        guess_mode = guess_mode or global_pool_conditions
+
+        # 3. Encode input prompt
+        text_encoder_lora_scale = (
+            self.cross_attention_kwargs.get("scale", None) if self.cross_attention_kwargs is not None else None
+        )
+        prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+            prompt,
+            device,
+            num_images_per_prompt,
+            self.do_classifier_free_guidance,
+            negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            lora_scale=text_encoder_lora_scale,
+            clip_skip=self.clip_skip,
+        )
+        # For classifier free guidance, we need to do two forward passes.
+        # Here we concatenate the unconditional and text embeddings into a single batch
+        # to avoid doing two forward passes
+        if self.do_classifier_free_guidance:
+            prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+                self.do_classifier_free_guidance,
+            )
+
+        # 4. Prepare image
+        if isinstance(controlnet, ControlNetModel):
+            qrcode = self.prepare_image(
+                image=qrcode,
+                width=width,
+                height=height,
+                batch_size=batch_size * num_images_per_prompt,
+                num_images_per_prompt=num_images_per_prompt,
+                device=device,
+                dtype=controlnet.dtype,
+                do_classifier_free_guidance=self.do_classifier_free_guidance,
+                guess_mode=guess_mode,
+            )
+            height, width = qrcode.shape[-2:]
+        elif isinstance(controlnet, MultiControlNetModel):
+            qrcodes = []
+
+            # Nested lists as ControlNet condition
+            if isinstance(qrcode[0], list):
+                # Transpose the nested image list
+                qrcode = [list(t) for t in zip(*qrcode)]
+
+            for qrcode_ in qrcode:
+                qrcode_ = self.prepare_image(
+                    image=qrcode_,
+                    width=width,
+                    height=height,
+                    batch_size=batch_size * num_images_per_prompt,
+                    num_images_per_prompt=num_images_per_prompt,
+                    device=device,
+                    dtype=controlnet.dtype,
+                    do_classifier_free_guidance=self.do_classifier_free_guidance,
+                    guess_mode=guess_mode,
+                )
+
+                qrcodes.append(qrcode_)
+
+            qrcode = qrcodes
+            height, width = qrcode[0].shape[-2:]
+        else:
+            assert False
+
+        # 5. Prepare timesteps
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler, num_inference_steps, device, timesteps, sigmas
+        )
+        self._num_timesteps = len(timesteps)
+
+        # 6. Prepare latent variables
+        num_channels_latents = self.unet.config.in_channels
+        latents = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 6.5 Optionally get Guidance Scale Embedding
+        timestep_cond = None
+        if self.unet.config.time_cond_proj_dim is not None:
+            guidance_scale_tensor = torch.tensor(self.guidance_scale - 1).repeat(batch_size * num_images_per_prompt)
+            timestep_cond = self.get_guidance_scale_embedding(
+                guidance_scale_tensor, embedding_dim=self.unet.config.time_cond_proj_dim
+            ).to(device=device, dtype=latents.dtype)
+
+        # 7. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+        extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+        # 7.1 Add image embeds for IP-Adapter
+        added_cond_kwargs = (
+            {"image_embeds": image_embeds}
+            if ip_adapter_image is not None or ip_adapter_image_embeds is not None
+            else None
+        )
+
+        # 7.2 Create tensor stating which controlnets to keep
+        controlnet_keep = []
+        for i in range(len(timesteps)):
+            keeps = [
+                1.0 - float(i / len(timesteps) < s or (i + 1) / len(timesteps) > e)
+                for s, e in zip(control_guidance_start, control_guidance_end)
+            ]
+            controlnet_keep.append(keeps[0] if isinstance(controlnet, ControlNetModel) else keeps)
+
+        # 8. Denoising loop
+        num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+        is_unet_compiled = is_compiled_module(self.unet)
+        is_controlnet_compiled = is_compiled_module(self.controlnet)
+        is_torch_higher_equal_2_1 = is_torch_version(">=", "2.1")
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            with torch.enable_grad():
+                for i, t in enumerate(timesteps):
+                    if self.interrupt:
+                        continue
+
+                    # Relevant thread:
+                    # https://dev-discuss.pytorch.org/t/cudagraphs-in-pytorch-2-0/1428
+                    if (is_unet_compiled and is_controlnet_compiled) and is_torch_higher_equal_2_1:
+                        torch._inductor.cudagraph_mark_step_begin()
+                    # expand the latents if we are doing classifier free guidance
+                    latents = latents.clone().detach().requires_grad_(True)
+                    latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
+                    latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+                    # controlnet(s) inference
+                    if guess_mode and self.do_classifier_free_guidance:
+                        # Infer ControlNet only for the conditional batch.
+                        control_model_input = latents
+                        control_model_input = self.scheduler.scale_model_input(control_model_input, t)
+                        controlnet_prompt_embeds = prompt_embeds.chunk(2)[1]
+                    else:
+                        control_model_input = latent_model_input
+                        controlnet_prompt_embeds = prompt_embeds
+
+                    if isinstance(controlnet_keep[i], list):
+                        cond_scale = [c * s for c, s in zip(controlnet_conditioning_scale, controlnet_keep[i])]
+                    else:
+                        controlnet_cond_scale = controlnet_conditioning_scale
+                        if isinstance(controlnet_cond_scale, list):
+                            controlnet_cond_scale = controlnet_cond_scale[0]
+                        cond_scale = controlnet_cond_scale * controlnet_keep[i]
+
+                    down_block_res_samples, mid_block_res_sample = self.controlnet(
+                        control_model_input,
+                        t,
+                        encoder_hidden_states=controlnet_prompt_embeds,
+                        controlnet_cond=qrcode,
+                        conditioning_scale=cond_scale,
+                        guess_mode=guess_mode,
+                        return_dict=False,
+                    )
+
+                    if guess_mode and self.do_classifier_free_guidance:
+                        # Inferred ControlNet only for the conditional batch.
+                        # To apply the output of ControlNet to both the unconditional and conditional batches,
+                        # add 0 to the unconditional batch to keep it unchanged.
+                        down_block_res_samples = [torch.cat([torch.zeros_like(d), d]) for d in down_block_res_samples]
+                        mid_block_res_sample = torch.cat([torch.zeros_like(mid_block_res_sample), mid_block_res_sample])
+
+                    # predict the noise residual
+                    noise_pred = self.unet(
+                        latent_model_input,
+                        t,
+                        encoder_hidden_states=prompt_embeds,
+                        timestep_cond=timestep_cond,
+                        cross_attention_kwargs=self.cross_attention_kwargs,
+                        down_block_additional_residuals=down_block_res_samples,
+                        mid_block_additional_residual=mid_block_res_sample,
+                        added_cond_kwargs=added_cond_kwargs,
+                        return_dict=False,
+                    )[0]
+
+                    # perform guidance
+                    if self.do_classifier_free_guidance:
+                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                        noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+                    # compute the original latents x_t -> x_0
+                    original_latents = self.scheduler.step(
+                        noise_pred,
+                        t,
+                        latents,
+                        **extra_step_kwargs,
+                        return_dict=True,
+                    ).pred_original_sample
+
+                    original_image = self.vae.decode(
+                        original_latents / self.vae.config.scaling_factor,
+                        return_dict=True,
+                    ).sample
+
+                    # compute the score of Scanninig Robust Perceptual Guidance (SRPG)
+                    score = self.srpg.compute_score(
+                        latents=latents,
+                        image=crop_padding(self.image_processor.denormalize(original_image), qrcode_padding),
+                        qrcode=crop_padding(image_binarize(qrcode[qrcode.size(0) // 2, None]), qrcode_padding),
+                        ref_image=crop_padding(ref_image, qrcode_padding),
+                    )
+
+                    timesteps_prev = t - self.scheduler.config.num_train_timesteps // self.scheduler.num_inference_steps
+                    alpha_prod_t = self.scheduler.alphas_cumprod[t]
+                    beta_prod_t = 1 - alpha_prod_t
+                    alpha_prod_t_prev = self.scheduler.alphas_cumprod[timesteps_prev] if timesteps_prev >= 0 else self.scheduler.final_alpha_cumprod
+                    beta_prod_t_prev = 1 - alpha_prod_t_prev
+
+                    noise_pred = noise_pred + (beta_prod_t ** 0.5) * score
+                    original_latents = (latents - (beta_prod_t ** 0.5) * noise_pred) / alpha_prod_t ** 0.5
+                    latents = (alpha_prod_t_prev ** 0.5) * original_latents + (beta_prod_t_prev ** 0.5) * noise_pred
+
+                    if callback_on_step_end is not None:
+                        callback_kwargs = {}
+                        for k in callback_on_step_end_tensor_inputs:
+                            callback_kwargs[k] = locals()[k]
+                        callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                        latents = callback_outputs.pop("latents", latents)
+                        prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+                        negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+
+                    # call the callback, if provided
+                    if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                        progress_bar.update()
+                        if callback is not None and i % callback_steps == 0:
+                            step_idx = i // getattr(self.scheduler, "order", 1)
+                            callback(step_idx, t, latents)
+
+                    if XLA_AVAILABLE:
+                        xm.mark_step()
+
+        # perform Scanning Robust Manifold Projected Gradient Descent (SR-MPGD)
+        if srmpgd_num_iteration is not None:
+            with torch.enable_grad():
+                latents = latents.clone().detach().requires_grad_(True)
+                optimizer = torch.optim.SGD([latents], lr=srmpgd_lr)
+
+                for i in trange(srmpgd_num_iteration):
+                    optimizer.zero_grad()
+                    original_image = self.vae.decode(latents / self.vae.config.scaling_factor,return_dict=False)[0]
+                    loss = self.srpg.compute_loss(
+                        image=crop_padding(self.image_processor.denormalize(original_image), qrcode_padding),
+                        qrcode=crop_padding(image_binarize(qrcode[qrcode.size(0) // 2, None]), qrcode_padding),
+                        ref_image=crop_padding(ref_image, qrcode_padding),
+                    )
+                    loss.backward()
+                    optimizer.step()
+
+        # If we do sequential model offloading, let's offload unet and controlnet
+        # manually for max memory savings
+        if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+            self.unet.to("cpu")
+            self.controlnet.to("cpu")
+            torch.cuda.empty_cache()
+
+        if not output_type == "latent":
+            image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False, generator=generator)[
+                0
+            ]
+            image, has_nsfw_concept = self.run_safety_checker(image, device, prompt_embeds.dtype)
+        else:
+            image = latents
+            has_nsfw_concept = None
+
+        if has_nsfw_concept is None:
+            do_denormalize = [True] * image.shape[0]
+        else:
+            do_denormalize = [not has_nsfw for has_nsfw in has_nsfw_concept]
+
+        image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image, has_nsfw_concept)
+
+        return StableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        qrcode: PipelineImageInput = None,
+        qrcode_module_size: int = 20,
+        qrcode_padding: int = 78,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        timesteps: List[int] = None,
+        sigmas: List[float] = None,
+        guidance_scale: float = 7.5,
+        negative_prompt: Optional[Union[str, List[str]]] = None,
+        num_images_per_prompt: Optional[int] = 1,
+        eta: float = 0.0,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        cross_attention_kwargs: Optional[Dict[str, Any]] = None,
+        controlnet_conditioning_scale: Union[float, List[float]] = 1.0,
+        guess_mode: bool = False,
+        control_guidance_start: Union[float, List[float]] = 0.0,
+        control_guidance_end: Union[float, List[float]] = 1.0,
+        scanning_robust_guidance_scale: int = 500,
+        perceptual_guidance_scale: int = 10,
+        clip_skip: Optional[int] = None,
+        srmpgd_num_iteration: Optional[int] = None,
+        srmpgd_lr: Optional[float] = 0.1,
+        callback_on_step_end: Optional[
+            Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+        ] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        **kwargs,
+    ):
+        stage1_output = self._run_stage1(
+            prompt=prompt,
+            qrcode=qrcode,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            timesteps=timesteps,
+            sigmas=sigmas,
+            guidance_scale=guidance_scale,
+            negative_prompt=negative_prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            eta=eta,
+            generator=generator,
+            latents=latents,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            ip_adapter_image=ip_adapter_image,
+            ip_adapter_image_embeds=ip_adapter_image_embeds,
+            output_type="pt",
+            return_dict=False,
+            cross_attention_kwargs=cross_attention_kwargs,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            guess_mode=guess_mode,
+            control_guidance_start=control_guidance_start,
+            control_guidance_end=control_guidance_end,
+            clip_skip=clip_skip,
+            callback_on_step_end=callback_on_step_end,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+        )
+        stage2_output = self._run_stage2(
+            prompt=prompt,
+            qrcode=qrcode,
+            qrcode_module_size=qrcode_module_size,
+            qrcode_padding=qrcode_padding,
+            ref_image=stage1_output.images,
+            height=height,
+            width=width,
+            num_inference_steps=num_inference_steps,
+            timesteps=timesteps,
+            sigmas=sigmas,
+            guidance_scale=guidance_scale,
+            negative_prompt=negative_prompt,
+            num_images_per_prompt=num_images_per_prompt,
+            eta=eta,
+            generator=generator,
+            latents=latents,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            ip_adapter_image=ip_adapter_image,
+            ip_adapter_image_embeds=ip_adapter_image_embeds,
+            output_type=output_type,
+            return_dict=return_dict,
+            cross_attention_kwargs=cross_attention_kwargs,
+            controlnet_conditioning_scale=controlnet_conditioning_scale,
+            guess_mode=guess_mode,
+            control_guidance_start=control_guidance_start,
+            control_guidance_end=control_guidance_end,
+            scanning_robust_guidance_scale=scanning_robust_guidance_scale,
+            perceptual_guidance_scale=perceptual_guidance_scale,
+            clip_skip=clip_skip,
+            srmpgd_num_iteration=srmpgd_num_iteration,
+            callback_on_step_end=callback_on_step_end,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+        )
+        return stage2_output

diffqrcoder/srpg.py

@@ -0,0 +1,33 @@
+import torch
+from torch import nn
+
+from diffqrcoder.losses import PerceptualLoss, ScanningRobustLoss
+
+
+
+GRADIENT_SCALE = 100
+
+
+class ScanningRobustPerceptualGuidance(nn.Module):
+    def __init__(
+        self,
+        module_size: int = 20,
+        scanning_robust_guidance_scale: int = 500,
+        perceptual_guidance_scale: int = 2,
+    ):
+        super().__init__()
+        self.module_size = module_size
+        self.scanning_robust_guidance_scale = scanning_robust_guidance_scale
+        self.perceptual_guidance_scale = perceptual_guidance_scale
+        self.scanning_robust_loss_fn = ScanningRobustLoss(module_size=module_size)
+        self.perceptual_loss_fn = PerceptualLoss()
+
+    def compute_loss(self, image: torch.Tensor, qrcode: torch.Tensor, ref_image: torch.Tensor) -> torch.Tensor:
+        return (
+            self.scanning_robust_guidance_scale * self.scanning_robust_loss_fn(image, qrcode) + \
+            self.perceptual_guidance_scale * self.perceptual_loss_fn(image, ref_image)
+        ) * GRADIENT_SCALE
+
+    def compute_score(self, latents: torch.Tensor, image: torch.Tensor, qrcode: torch.Tensor, ref_image: torch.Tensor) -> torch.Tensor:
+        loss = self.compute_loss(image, qrcode, ref_image)
+        return torch.autograd.grad(loss, latents)[0] / GRADIENT_SCALE

diffqrcoder/utils.py

@@ -0,0 +1,54 @@
+import numpy as np
+import torch
+from PIL import Image
+from torchvision.transforms.functional import pil_to_tensor
+
+
+def convert_pil_to_normalized_tensor(image: Image) -> torch.Tensor:
+    return pil_to_tensor(image).unsqueeze(0).float()
+
+
+def convert_normalized_tensor_to_np_image(image: torch.Tensor) -> np.array:
+    return image.clip(0, 1).squeeze(0).permute(1, 2, 0).detach().cpu().numpy()
+
+
+def add_position_pattern(
+    x: torch.Tensor,
+    y: torch.Tensor,
+    module_num: int,
+    module_size: int
+) -> torch.Tensor:
+
+    x[: 8 * module_size - 1, : 8 * module_size - 1, :] = \
+        y[: 8 * module_size - 1, : 8 * module_size - 1, :]
+
+    x[
+        (module_num - 8) * module_size + 1 : module_num * module_size,
+        : 8 * module_size - 1,
+        :
+    ] = y[
+        (module_num - 8) * module_size + 1 : module_num * module_size,
+        : 8 * module_size - 1,
+        :
+    ]
+
+    x[
+        : 8 * module_size - 1,
+        (module_num - 8) * module_size + 1 : module_num * module_size,
+        :
+    ] = y[
+        : 8 * module_size - 1,
+        (module_num - 8) * module_size + 1 : module_num * module_size,
+        :
+    ]
+
+    x[
+        (module_num - 9) * module_size : (module_num - 4) * module_size - 1,
+        (module_num - 9) * module_size : (module_num - 4) * module_size - 1,
+        :
+    ] = y[
+        (module_num - 9) * module_size : (module_num - 4) * module_size - 1,
+        (module_num - 9) * module_size : (module_num - 4) * module_size - 1,
+        :
+    ]
+    return x

diffqrcoder_wrapper.py

@@ -0,0 +1,121 @@
+# diffqrcoder_wrapper.py
+import torch
+from diffusers import ControlNetModel, DDIMScheduler
+from PIL import Image
+import qrcode
+
+from diffqrcoder import DiffQRCoderPipeline
+
+# ---- Defaults taken from run_diffqrcoder.py ---- #
+CONTROLNET_CKPT = "monster-labs/control_v1p_sd15_qrcode_monster"
+# Original used a direct file URL; we can keep that:
+PIPE_CKPT = (
+    "https://huggingface.co/fp16-guy/Cetus-Mix_Whalefall_fp16_cleaned/"
+    "resolve/main/cetusMix_Whalefall2_fp16.safetensors"
+)
+# You can also upload that file to the Space and use a local path.
+
+DEVICE = "cuda"  # ZeroGPU will give us a CUDA device during @spaces.GPU calls
+
+# Cache
+_controlnet = None
+_pipe = None
+
+
+def _make_qr_image(
+    data: str,
+    box_size: int = 20,   # aligns with qrcode_module_size default
+    border: int = 4,      # typical QR quiet zone in modules
+) -> Image.Image:
+    qr = qrcode.QRCode(
+        version=None,
+        error_correction=qrcode.constants.ERROR_CORRECT_H,
+        box_size=box_size,
+        border=border,
+    )
+    qr.add_data(data)
+    qr.make(fit=True)
+    img = qr.make_image(fill_color="black", back_color="white").convert("RGB")
+    return img
+
+
+def load_pipeline():
+    """
+    Lazily load ControlNet + DiffQRCoderPipeline.
+    Mirrors run_diffqrcoder.py, but only once.
+    """
+    global _controlnet, _pipe
+
+    if _pipe is not None:
+        return _pipe
+
+    # 1. ControlNet
+    if _controlnet is None:
+        _controlnet = ControlNetModel.from_pretrained(
+            CONTROLNET_CKPT,
+            torch_dtype=torch.float16,
+        )
+
+    # 2. DiffQRCoderPipeline (from single safetensors file)
+    pipe = DiffQRCoderPipeline.from_single_file(
+        PIPE_CKPT,
+        controlnet=_controlnet,
+        torch_dtype=torch.float16,
+        use_auth_token=True,   # uses the Space's HF token
+    )
+
+    # 3. Use DDIM scheduler as in original script
+    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+
+    # Don't call .to("cuda") yet; do it inside the @spaces.GPU function
+    _pipe = pipe
+    return _pipe
+
+
+def generate_qr_art(
+    url_or_text: str,
+    prompt: str,
+    neg_prompt: str = "easynegative",
+    num_inference_steps: int = 40,
+    qrcode_module_size: int = 20,
+    qrcode_padding: int = 78,
+    controlnet_conditioning_scale: float = 1.35,
+    scanning_robust_guidance_scale: float = 500.0,
+    perceptual_guidance_scale: float = 2.0,
+    srmpgd_num_iteration: int | None = None,
+    srmpgd_lr: float = 0.1,
+    seed: int = 1,
+) -> Image.Image:
+    """
+    Directly mirrors the call at the bottom of run_diffqrcoder.py,
+    but takes the QR content + prompt as arguments and returns a PIL image.
+    """
+    pipe = load_pipeline()
+
+    # ZeroGPU will ensure DEVICE exists as "cuda" when we call this
+    generator = torch.Generator(device=DEVICE).manual_seed(seed)
+
+    # Create QR image in-memory instead of loading from disk
+    qrcode_img = _make_qr_image(
+        data=url_or_text,
+        box_size=qrcode_module_size,  # roughly aligned
+        border=4,  # module-based border; padding param handles extra
+    )
+
+    pipe = pipe.to(DEVICE)
+
+    result = pipe(
+        prompt=prompt,
+        qrcode=qrcode_img,
+        qrcode_module_size=qrcode_module_size,
+        qrcode_padding=qrcode_padding,
+        negative_prompt=neg_prompt,
+        num_inference_steps=num_inference_steps,
+        generator=generator,
+        controlnet_conditioning_scale=controlnet_conditioning_scale,
+        scanning_robust_guidance_scale=scanning_robust_guidance_scale,
+        perceptual_guidance_scale=perceptual_guidance_scale,
+        srmpgd_num_iteration=srmpgd_num_iteration,
+        srmpgd_lr=srmpgd_lr,
+    )
+    return result.images[0]

requirements.txt

@@ -0,0 +1,12 @@
+gradio>=4.0.0
+spaces
+
+torch==2.1.2
+diffusers>=0.26.0
+transformers
+accelerate
+safetensors
+opencv-python
+numpy
+Pillow
+qrcode