Update README.md

README.md

@@ -134,12 +134,12 @@ pipe = FluxFillCFGPipeline.from_pretrained(
     torch_dtype=torch.bfloat16).to("cuda")
 ```
 
-### Object Removal
+### Image Extend with prompt
 ```python
-image = load_image('assets/image.png')
-mask = load_image('assets/mask_remove.png')
+image = load_image('assets/image2.png')
+mask = load_image('assets/mask_extend.png')
 image = pipe(
-    prompt='remove',  # using fix prompt in object removal
+    prompt='Deep in the forest, surronded by colorful flowers',
     negative_prompt="nsfw",
     image=image,
     mask_image=mask,
@@ -150,15 +150,15 @@ image = pipe(
     num_inference_steps=50,
     generator=torch.Generator("cpu").manual_seed(0)
 ).images[0]
-image.save(f"object_removal.jpg")
+image.save(f"image_extend_w_prompt.jpg")
 ```
 
-### Image Extend with prompt
+### Image Extend without prompt
 ```python
 image = load_image('assets/image2.png')
 mask = load_image('assets/mask_extend.png')
 image = pipe(
-    prompt='Deep in the forest, surronded by colorful flowers',
+    prompt='high-definition, perfect composition',  # using fix prompt in image extend wo prompt
     negative_prompt="nsfw",
     image=image,
     mask_image=mask,
@@ -169,15 +169,15 @@ image = pipe(
     num_inference_steps=50,
     generator=torch.Generator("cpu").manual_seed(0)
 ).images[0]
-image.save(f"image_extend_w_prompt.jpg")
+image.save(f"image_extend_wo_prompt.jpg")
 ```
 
-### Image Extend without prompt
+### Object Removal
 ```python
-image = load_image('assets/image2.png')
-mask = load_image('assets/mask_extend.png')
+image = load_image('assets/image.png')
+mask = load_image('assets/mask_remove.png')
 image = pipe(
-    prompt='high-definition, perfect composition',  # using fix prompt in image extend wo prompt
+    prompt='remove',  # using fix prompt in object removal
     negative_prompt="nsfw",
     image=image,
     mask_image=mask,
@@ -188,7 +188,55 @@ image = pipe(
     num_inference_steps=50,
     generator=torch.Generator("cpu").manual_seed(0)
 ).images[0]
-image.save(f"image_extend_wo_prompt.jpg")
+image.save(f"object_removal.jpg")
+```
+
+### Object Removal with Lora
+As the base model flux fill have undergone heavy SFT for object generation, the improvement on removal is not obvious. we release a lora for object removal separately and might be helpful for you.
+
+```python
+import torch
+from diffusers.utils import load_image
+from diffusers import FluxTransformer2DModel
+
+from src.pipeline_flux_fill_with_cfg import FluxFillCFGPipeline
+
+transformer_onereward = FluxTransformer2DModel.from_pretrained(
+    "bytedance-research/OneReward",
+    subfolder="flux.1-fill-dev-OneReward-transformer",
+    torch_dtype=torch.bfloat16
+)
+
+pipe = FluxFillCFGPipeline.from_pretrained(
+    "black-forest-labs/FLUX.1-Fill-dev", 
+    transformer=transformer_onereward,
+    torch_dtype=torch.bfloat16).to("cuda")
+
+pipe.load_lora_weights(
+    "bytedance-research/OneReward",
+    subfolder="flux.1-fill-dev-object-removal-lora",
+    weight_name="pytorch_lora_weights.safetensors",
+    adapter_name="object_removal_lora"
+)
+print("Loaded adapters:", pipe.get_list_adapters())  
+pipe.set_adapters(["object_removal_lora"], adapter_weights=[1.0])
+
+# Object Removal
+image = load_image('assets/image.png')
+mask = load_image('assets/mask_remove.png')
+image = pipe(
+    prompt='remove',  # using fix prompt in object removal
+    negative_prompt="nsfw",
+    image=image,
+    mask_image=mask,
+    height=image.height,
+    width=image.width,
+    guidance_scale=1.0,
+    true_cfg=4.0,
+    num_inference_steps=50,
+    generator=torch.Generator("cpu").manual_seed(0),
+).images[0]
+image.save(f"object_removal_lora.jpg")
 ```