Fix task_id caching for img2img

daras_ai_v2/stable_diffusion.py

@@ -236,6 +236,7 @@ def instruct_pix2pix(
     guidance_scale: float,
     image_guidance_scale: float,
     seed: int = 42,
+    task_id: str | None = None,
 ):
     return call_sd_multi(
         "diffusion.instruct_pix2pix",
@@ -256,6 +257,7 @@ def instruct_pix2pix(
             "image": images,
             "image_guidance_scale": image_guidance_scale,
         },
+        task_id=task_id,
     )
 
 
@@ -364,6 +366,7 @@ def img2img(
     negative_prompt: str = None,
     guidance_scale: float,
     seed: int = 42,
+    task_id: str | None = None,
 ):
     prompt_strength = prompt_strength or 0.7
     assert 0 <= prompt_strength <= 0.9, "Prompt Strength must be in range [0, 0.9]"
@@ -409,6 +412,7 @@ def img2img(
                     "image": [init_image],
                     "strength": prompt_strength,
                 },
+                task_id=task_id,
             )
     return [
         upload_file_from_bytes(f"gooey.ai - {prompt}.png", sd_img_bytes)

recipes/CompareText2Img.py

@@ -1,4 +1,3 @@
-import hashlib
 import typing
 
 from pydantic import BaseModel

recipes/Img2Img.py

@@ -9,6 +9,8 @@
 from daras_ai_v2.crypto import hash_together
 from daras_ai_v2.img_model_settings_widgets import img_model_settings
 from daras_ai_v2.loom_video_widget import youtube_video
+from daras_ai_v2.query_params import gooey_get_query_params
+from daras_ai_v2.query_params_util import extract_query_params
 from daras_ai_v2.stable_diffusion import (
     InpaintingModels,
     Img2ImgModels,
@@ -155,6 +157,8 @@ def run(self, state: dict) -> typing.Iterator[str | None]:
         if not self.request.user.disable_safety_checker:
             safety_checker(text=request.text_prompt, image=request.input_image)
 
+        _, run_id, uid = extract_query_params(gooey_get_query_params())
+        task_id = hash_together(run_id, uid) if run_id and uid else None
         if request.selected_model == Img2ImgModels.instruct_pix2pix.name:
             state["output_images"] = instruct_pix2pix(
                 prompt=request.text_prompt,
@@ -165,6 +169,7 @@ def run(self, state: dict) -> typing.Iterator[str | None]:
                 seed=request.seed,
                 images=[init_image],
                 image_guidance_scale=request.image_guidance_scale,
+                task_id=task_id,
             )
         elif request.selected_controlnet_model:
             state["output_images"] = controlnet(
@@ -178,6 +183,7 @@ def run(self, state: dict) -> typing.Iterator[str | None]:
                 guidance_scale=request.guidance_scale,
                 seed=request.seed,
                 controlnet_conditioning_scale=request.controlnet_conditioning_scale,
+                task_id=task_id,
             )
         else:
             state["output_images"] = img2img(
@@ -191,7 +197,7 @@ def run(self, state: dict) -> typing.Iterator[str | None]:
                 negative_prompt=request.negative_prompt,
                 guidance_scale=request.guidance_scale,
                 seed=request.seed,
-                task_id=hash_together(run_id, uid) if run_id and uid else None,
+                task_id=task_id,
             )
 
     def preview_description(self, state: dict) -> str: