Reference image support for QR code generation

daras_ai_v2/stable_diffusion.py

@@ -402,7 +402,7 @@ def controlnet(
     scheduler: str = None,
     prompt: str,
     num_outputs: int = 1,
-    init_image: str,
+    init_images: list[str] | str,
     num_inference_steps: int = 50,
     negative_prompt: str = None,
     guidance_scale: float = 7.5,
@@ -411,6 +411,8 @@ def controlnet(
 ):
     if isinstance(selected_controlnet_model, str):
         selected_controlnet_model = [selected_controlnet_model]
+    if isinstance(init_images, str):
+        init_images = [init_images] * len(selected_controlnet_model)
     prompt = add_prompt_prefix(prompt, selected_model)
     return call_sd_multi(
         "diffusion.controlnet",
@@ -432,7 +434,7 @@ def controlnet(
             "num_images_per_prompt": num_outputs,
             "num_inference_steps": num_inference_steps,
             "guidance_scale": guidance_scale,
-            "image": [init_image] * len(selected_controlnet_model),
+            "image": init_images,
             "controlnet_conditioning_scale": controlnet_conditioning_scale,
             # "strength": prompt_strength,
         },

recipes/Img2Img.py

@@ -165,7 +165,7 @@ def run(self, state: dict) -> typing.Iterator[str | None]:
                 selected_controlnet_model=request.selected_controlnet_model,
                 prompt=request.text_prompt,
                 num_outputs=request.num_outputs,
-                init_image=init_image,
+                init_images=init_image,
                 num_inference_steps=request.quality,
                 negative_prompt=request.negative_prompt,
                 guidance_scale=request.guidance_scale,

recipes/QRCodeGenerator.py

@@ -37,6 +37,7 @@
 from recipes.EmailFaceInpainting import get_photo_for_email
 from recipes.SocialLookupEmail import get_profile_for_email
 from url_shortener.models import ShortenedURL
+from daras_ai_v2.enum_selector_widget import enum_multiselect
 
 ATTEMPTS = 1
 DEFAULT_QR_CODE_META_IMG = "https://storage.googleapis.com/dara-c1b52.appspot.com/daras_ai/media/f09c8cfa-5393-11ee-a837-02420a000190/ai%20art%20qr%20codes1%201.png.png"
@@ -59,6 +60,12 @@ class QRCodeGeneratorPage(BasePage):
         obj_scale=0.65,
         obj_pos_x=0.5,
         obj_pos_y=0.5,
+        image_prompt_controlnet_models=[
+            ControlNetModels.sd_controlnet_canny.name,
+            ControlNetModels.sd_controlnet_depth.name,
+            ControlNetModels.sd_controlnet_tile.name,
+        ],
+        inspiration_strength=0.3,
     )
 
     def __init__(self, *args, **kwargs):
@@ -75,6 +82,11 @@ class RequestModel(BaseModel):
 
         text_prompt: str
         negative_prompt: str | None
+        image_prompt: str | None
+        image_prompt_controlnet_models: list[
+            typing.Literal[tuple(e.name for e in ControlNetModels)], ...
+        ] | None
+        inspiration_strength: float | None
 
         selected_model: typing.Literal[tuple(e.name for e in Text2ImgModels)] | None
         selected_controlnet_model: list[
@@ -129,6 +141,14 @@ def render_form_v2(self):
             key="text_prompt",
             placeholder="Bright sunshine coming through the cracks of a wet, cave wall of big rocks",
         )
+        st.file_uploader(
+            """
+            ### 🏞️ Reference Image [optional]
+            This image will be used as inspiration to blend with the QR Code.
+            """,
+            key="image_prompt",
+            accept=["image/*"],
+        )
 
         qr_code_source_key = "__qr_code_source"
         if qr_code_source_key not in st.session_state:
@@ -320,6 +340,30 @@ def render_settings(self):
             color=255,
         )
 
+        if st.session_state.get("image_prompt"):
+            st.write("---")
+            st.write(
+                """
+                ##### 🎨 Inspiration
+                Use this to control how the image prompt should influence the output.
+                """,
+                className="gui-input",
+            )
+            st.slider(
+                "Inspiration Strength",
+                min_value=0.0,
+                max_value=1.0,
+                step=0.05,
+                key="inspiration_strength",
+            )
+            enum_multiselect(
+                ControlNetModels,
+                label="Control Net Models",
+                key="image_prompt_controlnet_models",
+                checkboxes=False,
+                allow_none=False,
+            )
+
     def render_output(self):
         state = st.session_state
         self._render_outputs(state)
@@ -376,11 +420,22 @@ def run(self, state: dict) -> typing.Iterator[str | None]:
         state["raw_images"] = raw_images = []
 
         yield f"Running {Text2ImgModels[request.selected_model].value}..."
+        if isinstance(request.selected_controlnet_model, str):
+            request.selected_controlnet_model = [request.selected_controlnet_model]
+        init_images = [image] * len(request.selected_controlnet_model)
+        if request.image_prompt:
+            init_images += [request.image_prompt] * len(
+                request.image_prompt_controlnet_models
+            )
+            request.selected_controlnet_model += request.image_prompt_controlnet_models
+            request.controlnet_conditioning_scale += [
+                request.inspiration_strength
+            ] * len(request.image_prompt_controlnet_models)
         state["output_images"] = controlnet(
             selected_model=request.selected_model,
             selected_controlnet_model=request.selected_controlnet_model,
             prompt=request.text_prompt,
-            init_image=image,
+            init_images=init_images,
             num_outputs=request.num_outputs,
             num_inference_steps=request.quality,
             negative_prompt=request.negative_prompt,

scripts/run_all_diffusion.py

@@ -108,7 +108,7 @@
                     selected_controlnet_model=controlnet_model.name,
                     prompt=get_random_string(100, string.ascii_letters),
                     num_outputs=4,
-                    init_image=random_img,
+                    init_images=random_img,
                     num_inference_steps=1,
                     guidance_scale=7,
                 ),