app.py

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
import gradio_user_history as gr_user_history


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, profile: gr.OAuthProfile | None):
    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).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).images[0]  

    # First SD, then freeu
    images = [sd_image, freeu_image]

    gr_user_history.save_image(label=prompt + ' (SD)', image=sd_image, profile=profile, metadata={"prompt": prompt, "pipe": sd_options, "b1": 1.0, "b2": 1.0, "s1": 1.0, "s2": 1.0})
    gr_user_history.save_image(label=prompt + ' (FreeU)', image=freeu_image, profile=profile, metadata={"prompt": prompt, "pipe": "freeu", "b1": b1, "b2": b2, "s1": s1, "s2": s2})

    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():
            with gr.Accordion('FreeU Parameters (feel free to adjust these parameters based on your prompt): ', open=False):
                with gr.Row():             
                    sd_options = gr.Dropdown(["SD2.1", "SDXL"], label="SD options", value="SDXL", visible=True)
                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.3)
                    b2 = gr.Slider(label='b2: backbone factor of the second stage block of decoder',
                                            minimum=1,
                                            maximum=2.0,
                                            step=0.01,
                                            value=1.4)
                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")
        
    with gr.Group():
        with gr.Row():
            with gr.Accordion("Past generations", open=False):
                gr_user_history.render()
                
    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)