from diffusers import QwenImageEditPipeline
import torch
from PIL import Image

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

# Load trained LoRA weights for in-scene editing
pipeline.load_lora_weights("valiantcat/Qwen-Image-Edit-2509-Upscale2K", weight_name="qwen_image_edit_2509_upscale.safetensors")

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

# Define in-scene editing prompt
prompt = "Upscale this picture to 4K resolution."

# 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("edited_image.png")