import gradio as gr from PIL import Image import torch from diffusers import DiffusionPipeline from free_lunch_utils import register_free_upblock2d, register_free_crossattn_upblock2d model_id = "stabilityai/stable-diffusion-2-1" # model_id = "./stable-diffusion-2-1" pip_2_1 = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16) pip_2_1 = pip_2_1.to("cuda") model_id = "stabilityai/stable-diffusion-xl-base-1.0" pip_XL = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16) pip_XL = pip_XL.to("cuda") prompt_prev = None sd_options_prev = None seed_prev = None sd_image_prev = None def infer(prompt, sd_options, seed, b1, b2, s1, s2): global prompt_prev global sd_options_prev global seed_prev global sd_image_prev if sd_options == 'SD2.1': pip = pip_2_1 elif sd_options == 'SDXL': pip = pip_XL else: pip = pip_2_1 # pip = pip_2_1 run_baseline = False if prompt != prompt_prev or sd_options != sd_options_prev or seed != seed_prev: run_baseline = True prompt_prev = prompt sd_options_prev = sd_options seed_prev = seed if run_baseline: # register_free_upblock2d(pip, b1=1.0, b2=1.0, s1=1.0, s2=1.0) register_free_crossattn_upblock2d(pip, b1=1.0, b2=1.0, s1=1.0, s2=1.0) torch.manual_seed(seed) print("Generating SD:") sd_image = pip(prompt, num_inference_steps=25).images[0] sd_image_prev = sd_image else: sd_image = sd_image_prev # register_free_upblock2d(pip, b1=b1, b2=b2, s1=s1, s2=s1) register_free_crossattn_upblock2d(pip, b1=b1, b2=b2, s1=s1, s2=s1) torch.manual_seed(seed) print("Generating FreeU:") freeu_image = pip(prompt, num_inference_steps=25).images[0] # First SD, then freeu images = [sd_image, freeu_image] return images examples = [ [ "RAW photo, subject, 8k uhd, dslr, soft lighting, high quality, clearly face, a futuristic visage with cybernetic enhancements seamlessly integrated into human features", ], [ "Sculpt a life-sized animal using discarded plastic bottles and metal scraps, highlighting it's beauty, highly detailed, 8k", ], [ "A robot standing in the rain reading newspaper, rusty and worn down, in a dystopian cyberpunk street, photo-realistic , urbanpunk", ], [ "an outdoor full size sculpture using discarded car parts, highlighting it's beauty, highly detailed, 8k", ], [ "1955, moon landing, sci-fi, 8k, photorealistic, no atmosphere, earth in the sky, terraforming, style by Dean ellis", ], [ "a futuristic home , spaceship design,beautiful interior , high end design", ], [ "Hypnotic Maze, Fantasy Castle, Challenging Maze, Impossible Geometry, Mc Escher, Surreal Photography Within A Glass Sphere, Diorama, Beautiful Abundance, Medieval detailing , Digital Painting, Digital Illustration, Extreme Detail, Digital Art, 8k, Ultra Hd, Fantasy Art, Hyper Detailed, Hyperrealism, Elaborate, Vray, Unrea", ], [ "photo of half life combine standing outside city 17, glossy robot, rainy, rtx, octane, unreal", ], [ "new art : landscape into a Underground oasis in egypt. satara by johnny taylor, in the style of brushstroke-inmersive landscape, cinematic elegance, golden light, dark proportions, flowing brushwork, multilayered realism, --ar 61:128 --s 750 --v 5.2", ], [ "A horse galloping on the ocean", ], [ "a teddy bear walking in the snowstorm" ], [ "Campfire at night in a snowy forest with starry sky in the background." ], [ "a fantasy landscape, trending on artstation" ], [ "An astronaut flying in space, 4k, high resolution." ], [ "An astronaut is riding a horse in the space in a photorealistic style." ], [ "Turtle swimming in ocean." ], [ "A storm trooper vacuuming the beach." ], [ "Fireworks." ], [ "A fat rabbit wearing a purple robe walking through a fantasy landscape." ], [ "A koala bear playing piano in the forest." ], [ "An astronaut flying in space, 4k, high resolution." ], [ "Flying through fantasy landscapes, 4k, high resolution." ], [ "A small cabin on top of a snowy mountain in the style of Disney, artstation", ], [ "half human half cat, a human cat hybrid", ], [ "a drone flying over a snowy forest." ], ] css = """ h1 { text-align: center; } #component-0 { max-width: 730px; margin: auto; } """ block = gr.Blocks(css='style.css') options = ['SD2.1'] with block: gr.Markdown("# SD vs. FreeU") with gr.Group(): with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True): with gr.Column(): text = gr.Textbox( label="Enter your prompt", show_label=False, max_lines=1, placeholder="Enter your prompt", container=False, ) btn = gr.Button("Generate image", scale=0) with gr.Group(): with gr.Row(): sd_options = gr.Dropdown(["SD2.1", "SDXL"], label="SD options", value="SDXL", visible=True) with gr.Group(): with gr.Row(): with gr.Accordion('FreeU Parameters (feel free to adjust these parameters based on your prompt): ', open=False): with gr.Row(): b1 = gr.Slider(label='b1: backbone factor of the first stage block of decoder', minimum=1, maximum=2.0, step=0.01, value=1.4) b2 = gr.Slider(label='b2: backbone factor of the second stage block of decoder', minimum=1, maximum=2.0, step=0.01, value=1.6) with gr.Row(): s1 = gr.Slider(label='s1: skip factor of the first stage block of decoder', minimum=0, maximum=1, step=0.1, value=0.9) s2 = gr.Slider(label='s2: skip factor of the second stage block of decoder', minimum=0, maximum=1, step=0.1, value=0.2) seed = gr.Slider(label='seed', minimum=0, maximum=1000, step=1, value=42) with gr.Row(): with gr.Group(): # btn = gr.Button("Generate image", scale=0) with gr.Row(): with gr.Column() as c1: image_1 = gr.Image(interactive=False) image_1_label = gr.Markdown("SD") with gr.Group(): # btn = gr.Button("Generate image", scale=0) with gr.Row(): with gr.Column() as c2: image_2 = gr.Image(interactive=False) image_2_label = gr.Markdown("FreeU") ex = gr.Examples(examples=examples, fn=infer, inputs=[text, sd_options, seed, b1, b2, s1, s2], outputs=[image_1, image_2], cache_examples=False) ex.dataset.headers = [""] text.submit(infer, inputs=[text, sd_options, seed, b1, b2, s1, s2], outputs=[image_1, image_2]) btn.click(infer, inputs=[text, sd_options, seed, b1, b2, s1, s2], outputs=[image_1, image_2]) block.launch() # block.queue(default_enabled=False).launch(share=False)