Update app.py

app.py

@@ -18,7 +18,7 @@ from diffusers.utils import load_image
 MAX_SEED = np.iinfo(np.int32).max
 
 pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
-pipe.load_lora_weights("ovi054/Draw2Photo-Kontext-LoRA")
+pipe.load_lora_weights("ovi054/virtual-tryon-kontext-lora")
 pipe.fuse_lora()
 # optimize_pipeline_(pipe, image=Image.new("RGB", (512, 512)), prompt='prompt')
 
@@ -30,53 +30,53 @@ BASE_EXAMPLES = [os.path.join(EXAMPLES_DIR, "base", f) for f in sorted(os.listdi
 FACE_EXAMPLES = [os.path.join(EXAMPLES_DIR, "face", f) for f in sorted(os.listdir(os.path.join(EXAMPLES_DIR, "face")))]
 
 
-def add_overlay(base_img, overlay_img, margin=20):
-    """
-    Pastes an overlay image onto the top-right corner of a base image.
+# def add_overlay(base_img, overlay_img, margin=20):
+#     """
+#     Pastes an overlay image onto the top-right corner of a base image.
 
-    The overlay is resized to be 1/5th of the width of the base image,
-    maintaining its aspect ratio.
+#     The overlay is resized to be 1/5th of the width of the base image,
+#     maintaining its aspect ratio.
 
-    Args:
-        base_img (PIL.Image.Image): The main image.
-        overlay_img (PIL.Image.Image): The image to place on top.
-        margin (int, optional): The pixel margin from the top and right edges. Defaults to 20.
+#     Args:
+#         base_img (PIL.Image.Image): The main image.
+#         overlay_img (PIL.Image.Image): The image to place on top.
+#         margin (int, optional): The pixel margin from the top and right edges. Defaults to 20.
 
-    Returns:
-        PIL.Image.Image: The combined image.
-    """
-    if base_img is None or overlay_img is None:
-        return base_img
+#     Returns:
+#         PIL.Image.Image: The combined image.
+#     """
+#     if base_img is None or overlay_img is None:
+#         return base_img
     
-    base = base_img.convert("RGBA")
-    overlay = overlay_img.convert("RGBA")
+#     base = base_img.convert("RGBA")
+#     overlay = overlay_img.convert("RGBA")
     
-    # --- MODIFICATION ---
-    # Calculate the target width to be 1/5th of the base image's width
-    target_width = base.width // 5
+#     # --- MODIFICATION ---
+#     # Calculate the target width to be 1/5th of the base image's width
+#     target_width = base.width // 5
     
-    # Keep aspect ratio, resize overlay to the newly calculated target width
-    w, h = overlay.size
+#     # Keep aspect ratio, resize overlay to the newly calculated target width
+#     w, h = overlay.size
     
-    # Add a check to prevent division by zero if the overlay image has no width
-    if w == 0:
-        return base
+#     # Add a check to prevent division by zero if the overlay image has no width
+#     if w == 0:
+#         return base
         
-    new_height = int(h * (target_width / w))
-    overlay = overlay.resize((target_width, new_height), Image.LANCZOS)
+#     new_height = int(h * (target_width / w))
+#     overlay = overlay.resize((target_width, new_height), Image.LANCZOS)
 
-    # Position: top-right corner with a margin
-    x = base.width - overlay.width - margin
-    y = margin
+#     # Position: top-right corner with a margin
+#     x = base.width - overlay.width - margin
+#     y = margin
 
-    # Paste the resized overlay onto the base image using its alpha channel for transparency
-    base.paste(overlay, (x, y), overlay)
-    return base
+#     # Paste the resized overlay onto the base image using its alpha channel for transparency
+#     base.paste(overlay, (x, y), overlay)
+#     return base
 
 
 
 @spaces.GPU
-def infer(input_image, input_image_upload, overlay_image, prompt="make it real", seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, progress=gr.Progress(track_tqdm=True)):
+def infer(input_image_upload, prompt="wear it", seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, progress=gr.Progress(track_tqdm=True)):
     """
     Perform image editing using the FLUX.1 Kontext pipeline.
     
@@ -106,16 +106,7 @@ def infer(input_image, input_image_upload, overlay_image, prompt="make it real",
     # 1. Prioritize the uploaded image. If it exists, it becomes our main 'input_image'.
     if input_image_upload is not None:
         processed_input_image = input_image_upload
-    # 2. If no image was uploaded, check the drawing canvas.
-    elif isinstance(input_image, dict):
-        # Extract the actual image from the dictionary provided by gr.Paint
-        if "composite" in input_image and input_image["composite"] is not None:
-            processed_input_image = input_image["composite"]
-        elif "background" in input_image and input_image["background"] is not None:
-            processed_input_image = input_image["background"]
-        else:
-            # The canvas is empty, so there's no input image.
-            processed_input_image = None
+        
     else:
         # Fallback in case the input is neither from upload nor a valid canvas dict.
         processed_input_image = None
@@ -124,9 +115,6 @@ def infer(input_image, input_image_upload, overlay_image, prompt="make it real",
     
     # From this point on, 'processed_input_image' is either a PIL Image or None.
     if processed_input_image is not None:
-        if overlay_image is not None:
-            # Now this function is guaranteed to receive a PIL Image.
-            processed_input_image = add_overlay(processed_input_image, overlay_image)
             
         processed_input_image = processed_input_image.convert("RGB")
         image = pipe(
@@ -147,7 +135,7 @@ def infer(input_image, input_image_upload, overlay_image, prompt="make it real",
             generator=torch.Generator().manual_seed(seed),
         ).images[0]
         
-    return image, processed_input_image, seed, gr.Button(visible=False)
+    return image, seed, gr.Button(visible=False)
     
 @spaces.GPU
 def infer_example(input_image, prompt):
@@ -171,36 +159,26 @@ Turn drawing+face into a realistic photo with FLUX.1 Kontext [dev] + [Draw2Photo
         """)
         with gr.Row():
             with gr.Column():
-                gr.Markdown("Step 1.  Select/Upload/Draw a person ⬇️")
+                gr.Markdown("Step 1.  Select/Upload a model image + clothes overlay to try on ⬇️")
                 # input_image = gr.Image(label="Upload drawing", type="pil")
                 with gr.Row():
-                    with gr.Tabs() as tabs:
-                        with gr.TabItem("Upload"):
-                            input_image_upload = gr.Image(label="Upload drawing", type="pil")
-                            
-                        with gr.TabItem("Draw"):
-                            input_image = gr.Paint(
-                                    type="pil",
-                                    brush=gr.Brush(default_size=6, colors=["#000000"], color_mode="fixed"),
-                                    canvas_size = (1200,1200),
-                                    layers = False
-                                )
+                    input_image_upload = gr.Image(label="Upload drawing", type="pil")
                 gr.Examples(
                     examples=[[img] for img in BASE_EXAMPLES],
                     inputs=[input_image_upload],
                 )
 
-            with gr.Column():
-                gr.Markdown("Step 2.  Select/Upload a face photo ⬇️")
-                with gr.Row():
-                    overlay_image = gr.Image(label="Upload face photo", type="pil")
-                gr.Examples(
-                    examples=[[img] for img in FACE_EXAMPLES],
-                    inputs=[overlay_image],
-                )
+            # with gr.Column():
+            #     gr.Markdown("Step 2.  Select/Upload a face photo ⬇️")
+            #     with gr.Row():
+            #         overlay_image = gr.Image(label="Upload face photo", type="pil")
+            #     gr.Examples(
+            #         examples=[[img] for img in FACE_EXAMPLES],
+            #         inputs=[overlay_image],
+            #     )
                     
             with gr.Column():
-                gr.Markdown("Step 3.  Press “Run” to get results ⬇️")
+                gr.Markdown("Step 2.  Press “Run” to get results ⬇️")
                 with gr.Row():
                     run_button = gr.Button("Run")
                 with gr.Accordion("Advanced Settings", open=False):
@@ -208,7 +186,7 @@ Turn drawing+face into a realistic photo with FLUX.1 Kontext [dev] + [Draw2Photo
                     prompt = gr.Text(
                         label="Prompt",
                         max_lines=1,
-                        value = "make it real",
+                        value = "wear it",
                         placeholder="Enter your prompt for editing (e.g., 'Remove glasses', 'Add a hat')",
                         container=False,
                     )
@@ -258,8 +236,8 @@ Turn drawing+face into a realistic photo with FLUX.1 Kontext [dev] + [Draw2Photo
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn = infer,
-        inputs = [input_image, input_image_upload, overlay_image, prompt, seed, randomize_seed, guidance_scale, steps],
-        outputs = [result, result_input, seed, reuse_button]
+        inputs = [input_image_upload, prompt, seed, randomize_seed, guidance_scale, steps],
+        outputs = [result, seed, reuse_button]
     )
     # reuse_button.click(
     #     fn = lambda image: image,