Update app.py

app.py

@@ -1,26 +1,20 @@
 import base64
 from io import BytesIO
-
+import spaces
 import gradio as gr
 import torch
 from PIL import Image, ImageChops
-import spaces
-from diffsynth.pipelines.flux2_image import Flux2ImagePipeline, ModelConfig
+
+from diffusers import Flux2KleinPipeline
 
 from utils import process_source, process_reference, paste_back, binarize_mask
 
-pipe = Flux2ImagePipeline.from_pretrained(
-    torch_dtype=torch.bfloat16,
-    device="cuda",
-    model_configs=[
-        ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="text_encoder/*.safetensors"),
-        ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="transformer/*.safetensors"),
-        ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
-    ],
-    tokenizer_config=ModelConfig(model_id="black-forest-labs/FLUX.2-klein-base-4B", origin_file_pattern="tokenizer/"),
+pipe = Flux2KleinPipeline.from_pretrained(
+    "black-forest-labs/FLUX.2-klein-base-4B", torch_dtype=torch.bfloat16
 )
+pipe.to("cuda")
 
-pipe.load_lora(pipe.dit, ModelConfig(model_id="HuanJue/Easy-Insert", origin_file_pattern="easy-insert.safetensors"))
+pipe.load_lora_weights("LiXiY/Easy-Insert")
 
 PROMPT = "Replace the white mask of image1 with the content in image2. Preserving the background, lighting, and surrounding elements, maintain a seamless and natural result."
 
@@ -211,12 +205,15 @@ def load_example(idx):
 def load_ex1():
     return load_example(0)
 
+
 def load_ex2():
     return load_example(1)
 
+
 def load_ex3():
     return load_example(2)
 
+
 def load_ex4():
     return load_example(3)
 
@@ -291,7 +288,7 @@ def run_local(base, ref, base_state, ref_state, seed, num_inference_steps, cfg_s
                     text-align: center; font-size: 18px; box-shadow: 0 0 15px rgba(0,0,0,0.3);
                 ">
                     <p style="color: red; margin-bottom: 20px;">
-                        ⚠️ Please paint a mask on both the background image and the reference image first, or click an example!
+                        ⚠️ Please draw the mask on BOTH the background image and the reference image first, or click an example!
                     </p>
                     <button onclick="this.closest('.error-overlay').remove()"
                             style="padding: 8px 20px; cursor: pointer; border: none;
@@ -306,16 +303,16 @@ def run_local(base, ref, base_state, ref_state, seed, num_inference_steps, cfg_s
     background_image, _, crop_box, source_mask_cropped = process_source(pil_bg, pil_mask, CANVAS_W)
     ref_image = process_reference(pil_ref, pil_ref_mask, CANVAS_W)
 
+    generator = torch.Generator(device="cuda").manual_seed(int(seed))
     generated_image = pipe(
-        PROMPT,
-        edit_image=[background_image, ref_image],
-        seed=int(seed),
-        rand_device="cuda",
+        image=[background_image, ref_image],
+        prompt=PROMPT,
+        height=CANVAS_H,
+        width=CANVAS_W,
         num_inference_steps=int(num_inference_steps),
-        cfg_scale=float(cfg_scale),
-        height=CANVAS_W,
-        width=CANVAS_H,
-    )
+        guidance_scale=float(cfg_scale),
+        generator=generator,
+    ).images[0]
     result_img = paste_back(generated_image, pil_bg, crop_box, source_mask_cropped, feather=0)
 
     return result_img, gr.update(visible=False)
@@ -454,25 +451,21 @@ with gr.Blocks(css="""
 """) as demo:
 
     gr.Markdown(
-        "<h1 style='text-align: center;'>Reference-Based Object Insertion & Clothing Replacement</h1>"
-        "<h3 style='text-align: center;'>Insert an object from the reference image into the masked region of the background image, or replace the clothing in the masked region with the reference</h3>"
-        "<p style='text-align: center;'>"
-        "<a href='https://github.com/huan-yin/Easy-Insert' target='_blank' "
-        "style='color: #3b82f6; text-decoration: none; font-weight: 500;'>"
-        "GitHub Repo: huan-yin/Easy-Insert</a></p>"
+        "<h1 style='text-align: center;'>Reference-Based Object Insertion (FLUX.2-klein, diffusers)</h1>"
+        "<h3 style='text-align: center;'>Insert the object from the reference image into the masked region of the background image</h3>"
     )
     gr.Markdown(
         """
         **Instructions:**
-        1. Upload a background image and use the brush/eraser tools below it to mark or modify the insertion region (mask); upload a reference image and paint over the object you want to insert (reference mask).
-        2. Or click any thumbnail row in the "Examples" section below to automatically load the background image + mask + reference image + reference mask.
-        3. Click the "Generate" button; the result will be shown below.
+        1. Upload a background image, then use the brush/eraser tools below the image to mark or refine the insertion region (mask); upload a reference image and brush/erase over the object to be inserted (ref mask).
+        2. Or click any row of thumbnails in the "Examples" section below to automatically load background + mask + reference + ref mask.
+        3. Click the "Generate" button, and the result will be displayed below.
         """
     )
 
     with gr.Row(elem_classes="input-row"):
         base = gr.ImageEditor(
-            label="Background Image (paint the insertion region)",
+            label="Background Image (brush the insertion region)",
             type="pil",
             format="png",
             width=420,
@@ -489,7 +482,7 @@ with gr.Blocks(css="""
             eraser=gr.Eraser(default_size=30),
         )
         ref = gr.ImageEditor(
-            label="Reference Image (paint the object to insert)",
+            label="Reference Image (brush the object to insert)",
             type="pil",
             format="png",
             width=420,
@@ -507,7 +500,7 @@ with gr.Blocks(css="""
         )
 
     with gr.Row():
-        seed = gr.Number(label="Random Seed", value=1, precision=0)
+        seed = gr.Number(label="Seed", value=1, precision=0)
         num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=15)
         cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=10, step=0.1, value=4)
 
@@ -526,7 +519,7 @@ with gr.Blocks(css="""
     with gr.Row():
         error_dialog = gr.HTML(visible=False)
 
-    gr.HTML('<div class="ex-section-header"><span style="font-weight:700;font-size:16px;color:#374151;">Examples (click to load background image + mask + reference image + reference mask)</span></div>')
+    gr.HTML('<div class="ex-section-header"><span style="font-weight:700;font-size:16px;color:#374151;">Examples (click to load background + mask + reference + ref mask)</span></div>')
 
     base_state = gr.State(None)
     ref_state = gr.State(None)
@@ -555,7 +548,3 @@ with gr.Blocks(css="""
     )
 
 demo.launch(server_name="0.0.0.0", server_port=7860)
-
-
-
-