RiverZ/ICEdit

1'''2python scripts/gradio_demo.py 3'''4 5import sys6import os7workspace_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "icedit"))8 9if workspace_dir not in sys.path:10    sys.path.insert(0, workspace_dir)11    12from diffusers import FluxFillPipeline13import gradio as gr14import numpy as np15import torch16import argparse17import random 18import spaces19from PIL import Image20 21MAX_SEED = np.iinfo(np.int32).max22MAX_IMAGE_SIZE = 102423 24current_lora_scale = 1.025 26 27parser = argparse.ArgumentParser() 28parser.add_argument("--port", type=int, default=7860, help="Port for the Gradio app")29parser.add_argument("--output-dir", type=str, default="gradio_results", help="Directory to save the output image")30parser.add_argument("--flux-path", type=str, default='black-forest-labs/flux.1-fill-dev', help="Path to the model")31parser.add_argument("--lora-path", type=str, default='sanaka87/ICEdit-MoE-LoRA', help="Path to the LoRA weights")32parser.add_argument("--enable-model-cpu-offload", action="store_true", help="Enable CPU offloading for the model")33args = parser.parse_args()34 35pipe = FluxFillPipeline.from_pretrained(args.flux_path, torch_dtype=torch.bfloat16)36pipe.load_lora_weights(args.lora_path, adapter_name="icedit")37pipe.set_adapters("icedit", 1.0)38    39if args.enable_model_cpu_offload:40    pipe.enable_model_cpu_offload() 41else:42    pipe = pipe.to("cuda")43 44@spaces.GPU45def infer(edit_images, 46          prompt, 47          seed=666, 48          randomize_seed=False, 49          width=1024, 50          height=1024, 51          guidance_scale=50, 52          num_inference_steps=28, 53          lora_scale=1.0,54          progress=gr.Progress(track_tqdm=True)55):56    57 58    global current_lora_scale59    60    if lora_scale != current_lora_scale:61        print(f"\033[93m[INFO] LoRA scale changed from {current_lora_scale} to {lora_scale}, reloading LoRA weights\033[0m")62        pipe.set_adapters("icedit", lora_scale)63        current_lora_scale = lora_scale64        65    66    image = edit_images67        68    if image.size[0] != 512:69        print("\033[93m[WARNING] We can only deal with the case where the image's width is 512.\033[0m")70        new_width = 51271        scale = new_width / image.size[0]72        new_height = int(image.size[1] * scale)73        new_height = (new_height // 8) * 8  74        image = image.resize((new_width, new_height))75        print(f"\033[93m[WARNING] Resizing the image to {new_width} x {new_height}\033[0m")76        77    image = image.convert("RGB")78    width, height = image.size79    image = image.resize((512, int(512 * height / width)))80    combined_image = Image.new("RGB", (width * 2, height))81    combined_image.paste(image, (0, 0)) 82    mask_array = np.zeros((height, width * 2), dtype=np.uint8)83    mask_array[:, width:] = 255 84    mask = Image.fromarray(mask_array)85    instruction = f'A diptych with two side-by-side images of the same scene. On the right, the scene is exactly the same as on the left but {prompt}'86 87    if randomize_seed:88        seed = random.randint(0, MAX_SEED)89 90    output_image = pipe(91        prompt=instruction,92        image=combined_image,93        mask_image=mask,94        height=height,95        width=width*2,96        guidance_scale=guidance_scale,97        num_inference_steps=num_inference_steps,98        generator=torch.Generator().manual_seed(seed),99    ).images[0]100 101    w,h = output_image.size102    output_image = output_image.crop((w//2, 0, w, h))103 104    os.makedirs(args.output_dir, exist_ok=True)105        106    index = len(os.listdir(args.output_dir))107    output_image.save(f"{args.output_dir}/result_{index}.png")108    109    return (image, output_image), seed, lora_scale110    111# 新增的示例，将元组转换为列表112new_examples = [113    ['assets/girl_3.jpg', 'Make it looks like a watercolor painting.', 0, 0.5],114    ['assets/girl.png', 'Make her hair dark green and her clothes checked.', 42, 1.0],115    ['assets/boy.png', 'Change the sunglasses to a Christmas hat.', 27440001, 1.0],116    ['assets/kaori.jpg', 'Make it a sketch.', 329918865, 1.0]117]118 119css = """120#col-container {121    margin: 0 auto;122    max-width: 1000px;123}124"""125 126with gr.Blocks(css=css) as demo:127 128    with gr.Column(elem_id="col-container"):129        gr.Markdown(f"""# IC-Edit130**Image Editing is worth a single LoRA!** A demo for [IC-Edit](https://river-zhang.github.io/ICEdit-gh-pages/).131More **open-source**, with **lower costs**, **faster speed** (it takes about 9 seconds to process one image), and **powerful performance**.132For more details, check out our [Github Repository](https://github.com/River-Zhang/ICEdit) and [arxiv paper](https://arxiv.org/pdf/2504.20690). If our project resonates with you or proves useful, we'd be truly grateful if you could spare a moment to give it a star.133\n**👑 Feel free to share your results in this [Gallery](https://github.com/River-Zhang/ICEdit/discussions/21)!**134\n🔥 New feature: Try **different LoRA scale**!135""")136        with gr.Row():137            with gr.Column():138                edit_image = gr.Image(139                    label='Upload image for editing',140                    type='pil',141                    sources=["upload", "webcam"],142                    image_mode='RGB',143                    height=600144                )145                prompt = gr.Text(146                    label="Prompt",147                    show_label=False,148                    max_lines=1,149                    placeholder="Enter your prompt",150                    container=False,151                )152                run_button = gr.Button("Run")153            with gr.Column():154                result = gr.ImageSlider(label="Result", show_label=False)155                gr.Markdown("⚠️ If your edit didn't work as desired, **try again with another seed** ! <br> If you use our example, don't forget to uncheck the random seed option. Otherwise, it will still use a random seed.")156        with gr.Accordion("Advanced Settings", open=True):157 158            seed = gr.Slider(159                label="Seed",160                minimum=0,161                maximum=MAX_SEED,162                step=1,163                value=0,164            )165 166            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)167 168            with gr.Row():169 170                width = gr.Slider(171                    label="Width",172                    minimum=512,173                    maximum=MAX_IMAGE_SIZE,174                    step=32,175                    value=1024,176                    visible=False177                )178 179                height = gr.Slider(180                    label="Height",181                    minimum=512,182                    maximum=MAX_IMAGE_SIZE,183                    step=32,184                    value=1024,185                    visible=False186                )187 188            with gr.Row():189 190                guidance_scale = gr.Slider(191                    label="Guidance Scale",192                    minimum=1,193                    maximum=100,194                    step=0.5,195                    value=50,196                )197 198                num_inference_steps = gr.Slider(199                    label="Number of inference steps",200                    minimum=1,201                    maximum=50,202                    step=1,203                    value=28,204                )205 206            lora_scale = gr.Slider(207                label="LoRA Scale",208                minimum=0,209                maximum=1.0,210                step=0.01,211                value=1.0,212            )213        gr.Examples(214            examples=new_examples,215            inputs=[edit_image, prompt, seed, lora_scale],216            outputs=[result, seed, lora_scale],217            fn=infer,218            cache_examples=False219        )220 221    gr.on(222        triggers=[run_button.click, prompt.submit],223        fn=infer,224        inputs=[edit_image, prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, lora_scale],225        outputs=[result, seed, lora_scale]226    )227 228demo.launch(server_port=args.port)