CyStorm/instruct-pix2pix

1from __future__ import annotations2 3import math4import random5 6import gradio as gr7import torch8from PIL import Image, ImageOps9from diffusers import StableDiffusionInstructPix2PixPipeline10 11 12help_text = """13If you're not getting what you want, there may be a few reasons:141. Is the image not changing enough? Your Image CFG weight may be too high. This value dictates how similar the output should be to the input. It's possible your edit requires larger changes from the original image, and your Image CFG weight isn't allowing that. Alternatively, your Text CFG weight may be too low. This value dictates how much to listen to the text instruction. The default Image CFG of 1.5 and Text CFG of 7.5 are a good starting point, but aren't necessarily optimal for each edit. Try:15    * Decreasing the Image CFG weight, or16    * Increasing the Text CFG weight, or172. Conversely, is the image changing too much, such that the details in the original image aren't preserved? Try:18    * Increasing the Image CFG weight, or19    * Decreasing the Text CFG weight203. Try generating results with different random seeds by setting "Randomize Seed" and running generation multiple times. You can also try setting "Randomize CFG" to sample new Text CFG and Image CFG values each time.214. Rephrasing the instruction sometimes improves results (e.g., "turn him into a dog" vs. "make him a dog" vs. "as a dog").225. Increasing the number of steps sometimes improves results.236. Do faces look weird? The Stable Diffusion autoencoder has a hard time with faces that are small in the image. Try:24    * Cropping the image so the face takes up a larger portion of the frame.25"""26 27 28example_instructions = [29    "Make it a picasso painting",30    "as if it were by modigliani",31    "convert to a bronze statue",32    "Turn it into an anime.",33    "have it look like a graphic novel",34    "make him gain weight",35    "what would he look like bald?",36    "Have him smile",37    "Put him in a cocktail party.",38    "move him at the beach.",39    "add dramatic lighting",40    "Convert to black and white",41    "What if it were snowing?",42    "Give him a leather jacket",43    "Turn him into a cyborg!",44    "make him wear a beanie",45]46 47model_id = "timbrooks/instruct-pix2pix"48 49def main():50    pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16, safety_checker=None).to("cuda")51    example_image = Image.open("imgs/example.jpg").convert("RGB")52 53    def load_example(54        steps: int,55        randomize_seed: bool,56        seed: int,57        randomize_cfg: bool,58        text_cfg_scale: float,59        image_cfg_scale: float,60    ):61        example_instruction = random.choice(example_instructions)62        return [example_image, example_instruction] + generate(63            example_image,64            example_instruction,65            steps,66            randomize_seed,67            seed,68            randomize_cfg,69            text_cfg_scale,70            image_cfg_scale,71        )72 73    def generate(74        input_image: Image.Image,75        instruction: str,76        steps: int,77        randomize_seed: bool,78        seed: int,79        randomize_cfg: bool,80        text_cfg_scale: float,81        image_cfg_scale: float,82    ):83        seed = random.randint(0, 100000) if randomize_seed else seed84        text_cfg_scale = round(random.uniform(6.0, 9.0), ndigits=2) if randomize_cfg else text_cfg_scale85        image_cfg_scale = round(random.uniform(1.2, 1.8), ndigits=2) if randomize_cfg else image_cfg_scale86 87        width, height = input_image.size88        factor = 512 / max(width, height)89        factor = math.ceil(min(width, height) * factor / 64) * 64 / min(width, height)90        width = int((width * factor) // 64) * 6491        height = int((height * factor) // 64) * 6492        input_image = ImageOps.fit(input_image, (width, height), method=Image.Resampling.LANCZOS)93 94        if instruction == "":95            return [input_image, seed]96 97        generator = torch.manual_seed(seed)98        edited_image = pipe(99            instruction, image=input_image,100            guidance_scale=text_cfg_scale, image_guidance_scale=image_cfg_scale,101            num_inference_steps=steps, generator=generator,102        ).images[0]103        return [seed, text_cfg_scale, image_cfg_scale, edited_image]104 105    def reset():106        return [0, "Randomize Seed", 1371, "Fix CFG", 7.5, 1.5, None]107 108    with gr.Blocks() as demo:109        gr.HTML("""<h1 style="font-weight: 900; margin-bottom: 7px;">110   InstructPix2Pix: Learning to Follow Image Editing Instructions111</h1>112<p>For faster inference without waiting in queue, you may duplicate the space and upgrade to GPU in settings.113<br/>114<a href="https://huggingface.co/spaces/timbrooks/instruct-pix2pix?duplicate=true">115<img style="margin-top: 0em; margin-bottom: 0em" src="https://bit.ly/3gLdBN6" alt="Duplicate Space"></a>116<p/>""")117        with gr.Row():118            with gr.Column(scale=1, min_width=100):119                generate_button = gr.Button("Generate")120            with gr.Column(scale=1, min_width=100):121                load_button = gr.Button("Load Example")122            with gr.Column(scale=1, min_width=100):123                reset_button = gr.Button("Reset")124            with gr.Column(scale=3):125                instruction = gr.Textbox(lines=1, label="Edit Instruction", interactive=True)126 127        with gr.Row():128            input_image = gr.Image(label="Input Image", type="pil", interactive=True)129            edited_image = gr.Image(label=f"Edited Image", type="pil", interactive=False)130            input_image.style(height=512, width=512)131            edited_image.style(height=512, width=512)132 133        with gr.Row():134            steps = gr.Number(value=50, precision=0, label="Steps", interactive=True)135            randomize_seed = gr.Radio(136                ["Fix Seed", "Randomize Seed"],137                value="Randomize Seed",138                type="index",139                show_label=False,140                interactive=True,141            )142            seed = gr.Number(value=1371, precision=0, label="Seed", interactive=True)143            randomize_cfg = gr.Radio(144                ["Fix CFG", "Randomize CFG"],145                value="Fix CFG",146                type="index",147                show_label=False,148                interactive=True,149            )150            text_cfg_scale = gr.Number(value=7.5, label=f"Text CFG", interactive=True)151            image_cfg_scale = gr.Number(value=1.5, label=f"Image CFG", interactive=True)152 153        gr.Markdown(help_text)154 155        load_button.click(156            fn=load_example,157            inputs=[158                steps,159                randomize_seed,160                seed,161                randomize_cfg,162                text_cfg_scale,163                image_cfg_scale,164            ],165            outputs=[input_image, instruction, seed, text_cfg_scale, image_cfg_scale, edited_image],166        )167        generate_button.click(168            fn=generate,169            inputs=[170                input_image,171                instruction,172                steps,173                randomize_seed,174                seed,175                randomize_cfg,176                text_cfg_scale,177                image_cfg_scale,178            ],179            outputs=[seed, text_cfg_scale, image_cfg_scale, edited_image],180        )181        reset_button.click(182            fn=reset,183            inputs=[],184            outputs=[steps, randomize_seed, seed, randomize_cfg, text_cfg_scale, image_cfg_scale, edited_image],185        )186 187    demo.queue(concurrency_count=1)188    demo.launch(share=False)189 190 191if __name__ == "__main__":192    main()193