File size: 2,899 Bytes
be773bc 81376a7 be773bc 81376a7 dcfb6cf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
import gradio as gr
import random
import torch
import pathlib
from src.utils import concept_styles, loss_fn
from src.stable_diffusion import StableDiffusion
PROJECT_PATH = "."
CONCEPT_LIBS_PATH = f"{PROJECT_PATH}/concept_libs"
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
def generate(prompt, styles, gen_steps, loss_scale):
lossless_images, lossy_images = [], []
for style in styles:
concept_lib_path = f"{CONCEPT_LIBS_PATH}/{concept_styles[style]}"
concept_lib = pathlib.Path(concept_lib_path)
concept_embed = torch.load(concept_lib)
manual_seed = random.randint(0, 100)
diffusion = StableDiffusion(
device=DEVICE,
num_inference_steps=gen_steps,
manual_seed=manual_seed,
)
generated_image_lossless = diffusion.generate_image(
prompt=prompt,
loss_fn=loss_fn,
loss_scale=0,
concept_embed=concept_embed,
)
generated_image_lossy = diffusion.generate_image(
prompt=prompt,
loss_fn=loss_fn,
loss_scale=loss_scale,
concept_embed=concept_embed,
)
lossless_images.append((generated_image_lossless, style))
lossy_images.append((generated_image_lossy, style))
return {lossless_gallery: lossless_images, lossy_gallery: lossy_images}
with gr.Blocks() as app:
gr.Markdown("## ERA Session20 - Stable Diffusion: Generative Art with Guidance")
with gr.Row():
with gr.Column():
prompt_box = gr.Textbox(label="Prompt", interactive=True)
style_selector = gr.Dropdown(
choices=list(concept_styles.keys()),
value=list(concept_styles.keys())[0],
multiselect=True,
label="Select a Concept Style",
interactive=True,
)
gen_steps = gr.Slider(
minimum=10,
maximum=50,
value=30,
step=10,
label="Select Number of Steps",
interactive=True,
)
loss_scale = gr.Slider(
minimum=0,
maximum=32,
value=8,
step=8,
label="Select Guidance Scale",
interactive=True,
)
submit_btn = gr.Button(value="Generate")
with gr.Column():
lossless_gallery = gr.Gallery(
label="Generated Images without Guidance", show_label=True
)
lossy_gallery = gr.Gallery(
label="Generated Images with Guidance", show_label=True
)
submit_btn.click(
generate,
inputs=[prompt_box, style_selector, gen_steps, loss_scale],
outputs=[lossless_gallery, lossy_gallery],
)
app.launch()
|