Update pipeline.py

pipeline.py

@@ -75,49 +75,72 @@ def _safe_i2i_steps(strength: float, user_steps: int):
     steps_i2i = min(6, steps_i2i)                # keep fast
     return steps, steps_i2i, strength
 
-def run_search_and_generate(user_img: Image.Image, user_prompt: str,
+def run_search_and_generate(user_imgs, user_prompt: str,
                             k_retrieve=2, n_i2i=2, n_t2i=2,
                             strength_i2i=0.35, steps=1, gen_size=512, seed=42):
-    k_retrieve, n_i2i, n_t2i = _cap(k_retrieve), _cap(n_i2i), _cap(n_t2i)
+    # user_imgs: list of PIL images (1..4), some may be None
+    imgs = [im for im in (user_imgs or []) if im is not None]
+    if len(imgs) == 0:
+        raise ValueError("Please upload at least 1 image.")
+
     prompt = (user_prompt or "").strip()
     if not prompt:
         raise ValueError("Please enter a prompt (required for generation).")
 
-    q = embed_image(user_img)
+    k_retrieve, n_i2i, n_t2i = _cap(k_retrieve), _cap(n_i2i), _cap(n_t2i)
+
+    # --- embed each image and average embeddings ---
+    vecs = [embed_image(im) for im in imgs]
+    q = np.mean(np.stack(vecs, axis=0), axis=0)
+    q = (q / (np.linalg.norm(q) + 1e-12)).astype(np.float32)
+
+    # --- retrieval based on averaged embedding ---
     retrieved = retrieve(q, k_retrieve)
 
+    # --- choose init image for img2img (first provided) ---
+    init = imgs[0].convert("RGB").resize((gen_size, gen_size))
+
     steps_txt, steps_i2i, strength = _safe_i2i_steps(strength_i2i, steps)
-    init = user_img.convert("RGB").resize((gen_size, gen_size))
 
     gen_i2i, gen_t2i = [], []
+
     for i in range(n_i2i):
         g = torch.Generator(device).manual_seed(seed + 10*i)
         with torch.inference_mode():
             if device == "cuda":
                 with torch.autocast("cuda", dtype=torch.float16):
-                    gen_i2i.append(img2img(prompt=prompt, negative_prompt=NEG, image=init,
-                                           strength=strength, num_inference_steps=steps_i2i,
-                                           guidance_scale=0.0, generator=g).images[0])
+                    gen_i2i.append(img2img(
+                        prompt=prompt, negative_prompt=NEG, image=init,
+                        strength=strength, num_inference_steps=steps_i2i,
+                        guidance_scale=0.0, generator=g
+                    ).images[0])
             else:
-                gen_i2i.append(img2img(prompt=prompt, negative_prompt=NEG, image=init,
-                                       strength=strength, num_inference_steps=steps_i2i,
-                                       guidance_scale=0.0, generator=g).images[0])
+                gen_i2i.append(img2img(
+                    prompt=prompt, negative_prompt=NEG, image=init,
+                    strength=strength, num_inference_steps=steps_i2i,
+                    guidance_scale=0.0, generator=g
+                ).images[0])
 
     for i in range(n_t2i):
         g = torch.Generator(device).manual_seed(seed + 100 + 10*i)
         with torch.inference_mode():
             if device == "cuda":
                 with torch.autocast("cuda", dtype=torch.float16):
-                    gen_t2i.append(txt2img(prompt=prompt, negative_prompt=NEG,
-                                           num_inference_steps=steps_txt, guidance_scale=0.0,
-                                           height=gen_size, width=gen_size, generator=g).images[0])
+                    gen_t2i.append(txt2img(
+                        prompt=prompt, negative_prompt=NEG,
+                        num_inference_steps=steps_txt, guidance_scale=0.0,
+                        height=gen_size, width=gen_size, generator=g
+                    ).images[0])
             else:
-                gen_t2i.append(txt2img(prompt=prompt, negative_prompt=NEG,
-                                       num_inference_steps=steps_txt, guidance_scale=0.0,
-                                       height=gen_size, width=gen_size, generator=g).images[0])
+                gen_t2i.append(txt2img(
+                    prompt=prompt, negative_prompt=NEG,
+                    num_inference_steps=steps_txt, guidance_scale=0.0,
+                    height=gen_size, width=gen_size, generator=g
+                ).images[0])
 
     info = {
         "prompt": prompt,
+        "num_user_images": len(imgs),
         "k_retrieve": k_retrieve,
         "n_img2img": n_i2i,
         "n_txt2img": n_t2i,