Updated wrapper to use huggingface models.

diffqrcoder_wrapper.py

@@ -3,29 +3,29 @@ import torch
 from diffusers import ControlNetModel, DDIMScheduler
 from PIL import Image
 import qrcode
+from huggingface_hub import hf_hub_download
 
 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"
-# 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
+# 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"
+
+DEVICE = "cuda"
 
-# 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
+    box_size: int = 20,
+    border: int = 4,
 ) -> Image.Image:
     qr = qrcode.QRCode(
         version=None,
@@ -42,32 +42,42 @@ def _make_qr_image(
 def load_pipeline():
     """
     Lazily load ControlNet + DiffQRCoderPipeline.
-    Mirrors run_diffqrcoder.py, but only once.
+
+    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. ControlNet
+    # 1. Load ControlNet from its HF repo
     if _controlnet is None:
         _controlnet = ControlNetModel.from_pretrained(
             CONTROLNET_CKPT,
             torch_dtype=torch.float16,
         )
 
-    # 2. DiffQRCoderPipeline (from single safetensors file)
+    # 2. Download the base model safetensors from Hugging Face Hub
+    ckpt_path = hf_hub_download(
+        repo_id=PIPE_REPO_ID,
+        filename=PIPE_FILENAME,
+    )
+
+    # 3. Build DiffQRCoder pipeline from the local safetensors file
     pipe = DiffQRCoderPipeline.from_single_file(
-        PIPE_CKPT,
+        ckpt_path,
         controlnet=_controlnet,
         torch_dtype=torch.float16,
-        use_auth_token=True,   # uses the Space's HF token
+        use_auth_token=True,  # uses the Space's HF token
     )
 
-    # 3. Use DDIM scheduler as in original script
+    # 4. Same scheduler as in run_diffqrcoder.py
     pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
 
-    # Don't call .to("cuda") yet; do it inside the @spaces.GPU function
+    # NOTE: we call .to("cuda") inside the @spaces.GPU function so that
+    # it only happens when a GPU is actually attached.
     _pipe = pipe
     return _pipe
 
@@ -86,20 +96,14 @@ def generate_qr_art(
     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
+        box_size=qrcode_module_size,
+        border=4,
     )
 
     pipe = pipe.to(DEVICE)