|
import gradio as gr |
|
import tensorflow as tf |
|
from semdiffusers import SemanticEditPipeline |
|
|
|
device = 'gpu' if tf.config.list_physical_devices('GPU') else 'cpu' |
|
|
|
pipe = SemanticEditPipeline.from_pretrained( |
|
"runwayml/stable-diffusion-v1-5", |
|
).to(device) |
|
|
|
def infer(prompt, seed): |
|
gen = tf.random.Generator.from_seed(seed) |
|
|
|
out = pipe( |
|
prompt=prompt, |
|
generator=gen, |
|
num_images_per_prompt=1, |
|
guidance_scale=7 |
|
) |
|
|
|
images = out.images[0] |
|
|
|
out_edit = pipe( |
|
prompt=prompt, |
|
generator=gen, |
|
num_images_per_prompt=1, |
|
guidance_scale=7, |
|
editing_prompt=['male person', 'female person'], |
|
reverse_editing_direction=[True, False], |
|
edit_warmup_steps=[10, 10], |
|
edit_guidance_scale=[4, 4], |
|
edit_threshold=[0.95, 0.95], |
|
edit_momentum_scale=0.3, |
|
edit_mom_beta=0.6, |
|
edit_weights=[1, 1] |
|
) |
|
|
|
images_edited = out_edit.images[0] |
|
|
|
return [ |
|
(images, 'Stable Diffusion'), |
|
(images_edited, 'Fair Diffusion') |
|
] |
|
|
|
inputs = [ |
|
gr.inputs.Textbox(label='Prompt'), |
|
gr.inputs.Number(label='Seed', default=0, step=1) |
|
] |
|
|
|
outputs = gr.outputs.Image(label='Images', type='numpy', number=2) |
|
|
|
title = 'Semantic Edit Pipeline' |
|
description = 'Semantic Edit Pipeline implementation using SemDiffusers.' |
|
article = "<h3 style='text-align: center'><a href='https://github.com/crowsonkb/semdiffusers'>SemDiffusers</a></h3>" |
|
|
|
gr.Interface( |
|
infer, |
|
inputs, |
|
outputs, |
|
title=title, |
|
description=description, |
|
article=article, |
|
theme='compact' |
|
).launch(); |
|
|