Debugging

.gitignore

@@ -1,2 +1,10 @@
-__pycache__/*
-losses/__pycache__/*
+app.py
+__pycache__/diffqrcoder_wrapper.cpython-310.pyc
+diffqrcoder/__pycache__/__init__.cpython-310.pyc
+diffqrcoder/__pycache__/image_processor.cpython-310.pyc
+diffqrcoder/__pycache__/pipeline_diffqrcoder.cpython-310.pyc
+diffqrcoder/__pycache__/srpg.cpython-310.pyc
+diffqrcoder/losses/__pycache__/__init__.cpython-310.pyc
+diffqrcoder/losses/__pycache__/perceptual_loss.cpython-310.pyc
+diffqrcoder/losses/__pycache__/personalized_code_loss.cpython-310.pyc
+diffqrcoder/losses/__pycache__/scanning_robust_loss.cpython-310.pyc

app.py

@@ -4,6 +4,7 @@ import spaces
 from diffqrcoder_wrapper import generate_qr_art, load_pipeline
 import torch
 
+
 @spaces.GPU
 def infer(
     url_or_text: str,
@@ -14,22 +15,120 @@ def infer(
     perceptual_guidance_scale: float,
     srmpgd_iters: int,
 ):
-    # 🔹 Make sure pipeline is loaded *once* and then just moved to GPU
-    pipe = load_pipeline()
-
-    # move to GPU *here*, with an explicit non-blocking call
-    pipe = pipe.to("cuda")
-
-    srmpgd_num_iteration = None if srmpgd_iters == 0 else srmpgd_iters
-
-    img = generate_qr_art(
-        pipe,
-        url_or_text=url_or_text,
-        prompt=prompt,
-        num_inference_steps=num_inference_steps,
-        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,
+    try:
+        print("🔧 infer() starting")
+        print("CUDA available?", torch.cuda.is_available())
+        if torch.cuda.is_available():
+            print("CUDA device count:", torch.cuda.device_count())
+            print("Current device:", torch.cuda.current_device())
+            print("Device name:", torch.cuda.get_device_name(0))
+
+        pipe = load_pipeline()
+        print("✅ pipeline loaded on CPU")
+
+        # Attach to GPU in ZeroGPU context
+        pipe = pipe.to("cuda")
+        print("✅ pipeline moved to CUDA")
+
+        srmpgd_num_iteration = None if srmpgd_iters == 0 else srmpgd_iters
+        print(
+            f"Params → steps={num_inference_steps}, "
+            f"ctrl={controlnet_scale}, srg={scanning_robust_guidance_scale}, "
+            f"pg={perceptual_guidance_scale}, iters={srmpgd_num_iteration}"
+        )
+
+        img = generate_qr_art(
+            pipe,
+            url_or_text=url_or_text,
+            prompt=prompt,
+            num_inference_steps=num_inference_steps,
+            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,
+        )
+
+        print("✅ generation complete")
+        return img
+
+    except Exception as e:
+        print("❌ Error in infer():", repr(e))
+        raise
+
+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)) 🚀
+        """
     )
-    return img
+
+    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.launch()

diffqrcoder_wrapper.py

@@ -9,10 +9,8 @@ from diffqrcoder import DiffQRCoderPipeline
 
 # ---- Defaults taken from run_diffqrcoder.py ---- #
 
-# ControlNet is already a proper HF repo id:
 CONTROLNET_CKPT = "monster-labs/control_v1p_sd15_qrcode_monster"
 
-# For the base SD model (Cetus-Mix), use repo + filename rather than raw URL
 PIPE_REPO_ID = "fp16-guy/Cetus-Mix_Whalefall_fp16_cleaned"
 PIPE_FILENAME = "cetusMix_Whalefall2_fp16.safetensors"
 
@@ -42,30 +40,30 @@ def _make_qr_image(
 def load_pipeline():
     """
     Lazily load ControlNet + DiffQRCoderPipeline.
-
-    This now:
-    - pulls the ControlNet weights from HF by repo id
-    - downloads the Cetus-Mix safetensors file via hf_hub_download
     """
     global _controlnet, _pipe
 
     if _pipe is not None:
         return _pipe
 
-    # 1. Load ControlNet from its HF repo
+    print("🔧 Loading ControlNet...")
     if _controlnet is None:
         _controlnet = ControlNetModel.from_pretrained(
             CONTROLNET_CKPT,
             torch_dtype=torch.float16,
         )
+    print("✅ ControlNet loaded.")
 
-    # 2. Download the base model safetensors from Hugging Face Hub
+    print("🔧 Downloading base model checkpoint from Hub...")
     ckpt_path = hf_hub_download(
         repo_id=PIPE_REPO_ID,
         filename=PIPE_FILENAME,
+        local_dir="models",
+        local_dir_use_symlinks=False,
     )
+    print("✅ Base model checkpoint at:", ckpt_path)
 
-    # 3. Build DiffQRCoder pipeline from the local safetensors file
+    print("🔧 Building DiffQRCoderPipeline from checkpoint...")
     pipe = DiffQRCoderPipeline.from_single_file(
         ckpt_path,
         controlnet=_controlnet,
@@ -73,40 +71,47 @@ def load_pipeline():
         use_auth_token=True,  # uses the Space's HF token
     )
 
-    # 4. Same scheduler as in run_diffqrcoder.py
     pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
 
-    # NOTE: we call .to("cuda") inside the @spaces.GPU function so that
-    # it only happens when a GPU is actually attached.
+    # Memory helpers – cheaper attention
+    try:
+        pipe.enable_attention_slicing()
+        # Optional: pipe.enable_xformers_memory_efficient_attention()
+    except Exception as e:
+        print("⚠️ Could not enable attention optimizations:", repr(e))
+
+    print("✅ Pipeline constructed on CPU.")
     _pipe = pipe
     return _pipe
 
 
 def generate_qr_art(
-    pipe: DiffQRCoderPipeline, 
+    pipe: DiffQRCoderPipeline,
     url_or_text: str,
     prompt: str,
     neg_prompt: str = "easynegative",
-    num_inference_steps: int = 40,
+    num_inference_steps: int = 20,          # gentler default
     qrcode_module_size: int = 20,
     qrcode_padding: int = 78,
     controlnet_conditioning_scale: float = 1.35,
-    scanning_robust_guidance_scale: float = 500.0,
+    scanning_robust_guidance_scale: float = 300.0,  # softer default
     perceptual_guidance_scale: float = 2.0,
-    srmpgd_num_iteration: int | None = None,
+    srmpgd_num_iteration: int | None = 0,   # 0 = disable SR-MPGD by default
     srmpgd_lr: float = 0.1,
     seed: int = 1,
 ) -> Image.Image:
+    assert pipe is not None, "Pipeline must be loaded before calling generate_qr_art"
+
+    print("✨ generate_qr_art() starting...")
+    generator = torch.Generator(device=DEVICE).manual_seed(seed)
 
-    generator = torch.Generator(device="cuda").manual_seed(seed)
     qrcode_img = _make_qr_image(
         data=url_or_text,
         box_size=qrcode_module_size,
         border=4,
     )
 
-    # pipe = pipe.to(DEVICE)
-
+    print("✨ Starting DiffQRCoder forward pass...")
     result = pipe(
         prompt=prompt,
         qrcode=qrcode_img,
@@ -121,4 +126,5 @@ def generate_qr_art(
         srmpgd_num_iteration=srmpgd_num_iteration,
         srmpgd_lr=srmpgd_lr,
     )
+    print("✅ DiffQRCoder forward pass finished.")
     return result.images[0]