Spaces:
Sleeping
Sleeping
import os | |
import time | |
from datetime import datetime, timezone, timedelta | |
import spaces | |
import torch | |
import numpy as np | |
import gradio as gr | |
from huggingface_hub import hf_hub_download | |
from utils import preprocess_img, postprocess_img, load_model_without_module | |
from vgg.vgg19 import VGG_19 | |
from u2net.model import U2Net | |
from inference import inference | |
if torch.cuda.is_available(): device = 'cuda' | |
elif torch.backends.mps.is_available(): device = 'mps' | |
else: device = 'cpu' | |
print('Device:', device) | |
if device == 'cuda': print('Name:', torch.cuda.get_device_name()) | |
# load models | |
model = VGG_19().to(device).eval() | |
for param in model.parameters(): | |
param.requires_grad = False | |
sod_model = U2Net().to(device).eval() | |
load_model_without_module( | |
sod_model, | |
hf_hub_download(repo_id='jamino30/u2net-saliency', filename='u2net-duts-msra.safetensors'), | |
device=device | |
) | |
style_files = os.listdir('./style_images') | |
style_options = { | |
'Starry Night': './style_images/Starry_Night.jpg', | |
'Starry Night (v2)': './style_images/Starry_Night_v2.jpg', | |
'Scream': './style_images/Scream.jpg', | |
'Great Wave': './style_images/Great_Wave.jpg', | |
'Oil Painting': './style_images/Oil_Painting.jpg', | |
'Watercolor': './style_images/Watercolor.jpg', | |
'Mosaic': './style_images/Mosaic.jpg', | |
'Lego Bricks': './style_images/Lego_Bricks.jpg', | |
'Bokeh': './style_images/Bokeh.jpg', | |
} | |
lrs = np.linspace(0.015, 0.075, 3).tolist() | |
img_size = 512 | |
cached_style_features = { | |
style_name: model(preprocess_img(style_img_path, img_size)[0].to(device)) | |
for style_name, style_img_path in style_options.items() | |
} | |
def run(content_image, style_name, style_strength=len(lrs), apply_to_background=False): | |
yield None | |
content_img, original_size = preprocess_img(content_image, img_size) | |
content_img_normalized, _ = preprocess_img(content_image, img_size, normalize=True) | |
content_img, content_img_normalized = content_img.to(device), content_img_normalized.to(device) | |
style_features = cached_style_features[style_name] | |
print('-'*30) | |
print(datetime.now(timezone.utc) - timedelta(hours=5)) # EST | |
st = time.time() | |
generated_img = inference( | |
model=model, | |
sod_model=sod_model, | |
content_image=content_img, | |
content_image_norm=content_img_normalized, | |
style_features=style_features, | |
lr=lrs[style_strength-1], | |
apply_to_background=apply_to_background, | |
) | |
print(f'{time.time()-st:.2f}s') | |
yield postprocess_img(generated_img, original_size) | |
css = """ | |
#container { | |
margin: 0 auto; | |
max-width: 1200px; | |
} | |
""" | |
with gr.Blocks(css=css) as demo: | |
gr.HTML("<h1 style='text-align: center; padding: 10px'>🖼️ Neural Style Transfer w/ Salient Region Preservation") | |
with gr.Row(elem_id='container'): | |
with gr.Column(): | |
with gr.Group(): | |
content_image = gr.Image(label='Content', type='pil', sources=['upload', 'webcam', 'clipboard'], format='jpg', show_download_button=False) | |
with gr.Group(): | |
style_dropdown = gr.Radio(choices=list(style_options.keys()), label='Style', value='Starry Night', type='value') | |
style_strength_slider = gr.Slider(label='Style Strength', minimum=1, maximum=len(lrs), step=1, value=len(lrs)) | |
apply_to_background_checkbox = gr.Checkbox(label='Apply style transfer exclusively to the background', value=False) | |
submit_button = gr.Button('Submit', variant='primary') | |
examples = gr.Examples( | |
examples=[ | |
['./content_images/GoldenRetriever.jpg', 'Great Wave'], | |
['./content_images/CameraGirl.jpg', 'Bokeh'] | |
], | |
inputs=[content_image, style_dropdown] | |
) | |
with gr.Column(): | |
output_image = gr.Image(label='Output', type='pil', interactive=False, show_download_button=False) | |
download_button = gr.DownloadButton(label='Download Image', visible=False) | |
def save_image(img): | |
filename = 'generated.jpg' | |
img.save(filename) | |
return filename | |
submit_button.click( | |
fn=lambda: gr.update(visible=False), | |
outputs=download_button | |
) | |
submit_button.click( | |
fn=run, | |
inputs=[content_image, style_dropdown, style_strength_slider, apply_to_background_checkbox], | |
outputs=output_image | |
).then( | |
fn=save_image, | |
inputs=output_image, | |
outputs=download_button | |
).then( | |
fn=lambda: gr.update(visible=True), | |
outputs=download_button | |
) | |
demo.queue = False | |
demo.config['queue'] = False | |
demo.launch(show_api=False) | |