from diffusers import QwenImageEditPipeline
import torch
from PIL import Image

# Load the pipeline
pipeline = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit")
pipeline.to(torch.bfloat16)
pipeline.to("cuda")

# Load trained LoRA weights for in-scene editing
pipeline.load_lora_weights("valiantcat/Qwen-Image-Edit-Remover-General-LoRA",weight_name="qwen-edit-remover.safetensors")

# Load input image
image = Image.open("./result/test.png").convert("RGB")

# Define in-scene editing prompt
prompt = "移除猫猫"

# Generate edited image with enhanced scene understanding
inputs = {
    "image": image,
    "prompt": prompt,
    "generator": torch.manual_seed(12345),
    "true_cfg_scale": 4.0,
    "negative_prompt": " ",
    "num_inference_steps": 50,
}

with torch.inference_mode():
    output = pipeline(**inputs)
    output_image = output.images[0]
    output_image.save("restlt.png")