Allow uploading a background image

app.py

@@ -43,38 +43,38 @@ pipe.enable_xformers_memory_efficient_attention()
 # pipe.unet = torch.nn.DataParallel(pipe.unet)
 
 
-def resize_for_condition_image(input_image: Image.Image, resolution: int, canvas_width: int = None, canvas_height: int = None):
+def resize_for_condition_image(input_image: Image.Image, resolution: int, background_image: Image.Image = None, canvas_width: int = None, canvas_height: int = None):
     input_image = input_image.convert("RGBA")
     W, H = input_image.size
 
-    # Create a blank canvas with the specified dimensions and 50% transparent white background
-    if canvas_width is None:
-        canvas_width = W
-    if canvas_height is None:
-        canvas_height = H
-
-    canvas = Image.new("RGBA", (canvas_width, canvas_height), (255, 255, 255, 0))
+    # Use the uploaded background image or create a blank canvas
+    if background_image:
+        background_image = background_image.convert("RGBA")
+        canvas = background_image.resize((canvas_width or W, canvas_height or H), Image.LANCZOS)
+    else:
+        canvas = Image.new("RGBA", (canvas_width or W, canvas_height or H), (255, 255, 255, 0))
 
     # Determine the relative size of the QR code based on the canvas dimensions
     qr_scale_ratio = 768 / min(1024, 768)  # Base ratio (relative to 1024x768)
-    qr_target_size = int(min(canvas_width, canvas_height) * qr_scale_ratio)
+    qr_target_size = int(min(canvas.size) * qr_scale_ratio)
 
     # Resize the QR code to maintain its relative size
     input_image = input_image.resize((qr_target_size, qr_target_size), resample=Image.LANCZOS)
     W, H = input_image.size
 
-    # Paste the resized QR code 2/3rds of the way over the canvas
-    qr_x = int(canvas_width * (2 / 3)) - (W // 2)  # Adjust x-coordinate to 2/3 of the canvas
-    qr_y = (canvas_height - H) // 2  # Center the QR code vertically
+    # Paste the resized QR code onto the background
+    qr_x = int(canvas.size[0] * (2 / 3)) - (W // 2)  # Adjust x-coordinate to 2/3 of the canvas
+    qr_y = (canvas.size[1] - H) // 2  # Center the QR code vertically
 
     # Use the alpha channel of the input_image as the mask
-    canvas.paste(input_image, (qr_x, qr_y), mask=input_image.split()[3])  # Alpha channel as mask
+    canvas.paste(input_image, (qr_x, qr_y), mask=input_image.split()[3])
 
     return canvas
 
 
 
 
+
 SAMPLER_MAP = {
     "DPM++ Karras SDE": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True, algorithm_type="sde-dpmsolver++"),
     "DPM++ Karras": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True),
@@ -95,6 +95,7 @@ def inference(
     seed: int = -1,
     init_image: Image.Image | None = None,
     qrcode_image: Image.Image | None = None,
+    background_image: Image.Image | None = None,  # New input
     use_qr_code_as_init_image=True,
     sampler="DPM++ Karras SDE",
     width: int = 768,
@@ -126,15 +127,15 @@ def inference(
         # Add 50% transparency to the white background
         data = qrcode_image.getdata()
         new_data = [
-            (item[0], item[1], item[2], 26) if item[:3] == (255, 255, 255) else item
+            (item[0], item[1], item[2], 128) if item[:3] == (255, 255, 255) else item
             for item in data
         ]
         qrcode_image.putdata(new_data)
 
-        qrcode_image = resize_for_condition_image(qrcode_image, 768, width, height)
+        qrcode_image = resize_for_condition_image(qrcode_image, 768, background_image, width, height)
     else:
         print("Using QR Code Image")
-        qrcode_image = resize_for_condition_image(qrcode_image, 768, width, height)
+        qrcode_image = resize_for_condition_image(qrcode_image, 768, background_image, width, height)
 
     # hack due to gradio examples
     init_image = qrcode_image
@@ -155,6 +156,7 @@ def inference(
     return out.images[0]  # type: ignore
 
 
+
 with gr.Blocks() as blocks:
     gr.Markdown(
         """
@@ -204,6 +206,12 @@ model: https://huggingface.co/DionTimmer/controlnet_qrcode-control_v1p_sd15
                 init_image = gr.Image(label="Init Image (Optional). Leave blank to generate image with SD 2.1", type="pil")
 
 
+            with gr.Accordion(label="Background Image (Optional)", open=False):
+                background_image = gr.Image(
+                    label="Background Image (Optional). Leave blank for default background",
+                    type="pil",
+                )
+
             with gr.Accordion(
                 label="Params: The generated QR Code functionality is largely influenced by the parameters detailed below",
                 open=True,
@@ -253,26 +261,26 @@ model: https://huggingface.co/DionTimmer/controlnet_qrcode-control_v1p_sd15
                 run_btn = gr.Button("Run")
         with gr.Column():
             result_image = gr.Image(label="Result Image")
-    run_btn.click(
-        inference,
-        inputs=[
-            qr_code_content,
-            prompt,
-            negative_prompt,
-            guidance_scale,
-            controlnet_conditioning_scale,
-            strength,
-            seed,
-            init_image,
-            qr_code_image,
-            use_qr_code_as_init_image,
-            sampler,
-            width,  # add width input
-            height,  # add height input
-        ],
-        outputs=[result_image],
-        concurrency_limit=1
-    )
+        run_btn.click(
+            inference,
+            inputs=[
+                qr_code_content,
+                prompt,
+                negative_prompt,
+                guidance_scale,
+                controlnet_conditioning_scale,
+                strength,
+                seed,
+                init_image,
+                qr_code_image,
+                use_qr_code_as_init_image,
+                sampler,
+                width,  # add width input
+                height,  # add height input
+            ],
+            outputs=[result_image],
+            concurrency_limit=1
+        )
 
     gr.Examples(
         examples=[