Upload folder using huggingface_hub

app.py

@@ -4,13 +4,11 @@ from datetime import datetime, timezone, timedelta
 
 import spaces
 import torch
-import torch.optim as optim
 import numpy as np
 import gradio as gr
-from safetensors.torch import load_file
 from huggingface_hub import hf_hub_download
 
-from utils import preprocess_img, postprocess_img
+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
@@ -20,22 +18,20 @@ elif torch.backends.mps.is_available(): device = 'mps'
 else: device = 'cpu'
 print('Device:', device)
 if device == 'cuda': print('Name:', torch.cuda.get_device_name())
-
-def load_model_without_module(model, model_path):
-    state_dict = load_file(model_path, device=device)
-    new_state_dict = {}
-    for k, v in state_dict.items():
-        name = k[7:] if k.startswith('module.') else k
-        new_state_dict[name] = v
-    model.load_state_dict(new_state_dict)
    
 # load models 
 model = VGG_19().to(device).eval()
 for param in model.parameters():
     param.requires_grad = False
 sod_model = U2Net().to(device).eval()
-local_model_path = hf_hub_download(repo_id='jamino30/u2net-saliency', filename='u2net-duts-msra.safetensors')
-load_model_without_module(sod_model, local_model_path)
+load_model_without_module(
+    sod_model, 
+    hf_hub_download(repo_id='jamino30/u2net-saliency', filename='u2net-duts-msra.safetensors'),
+    device=device
+)
+
+model = torch.jit.script(model)
+sod_model = torch.jit.script(sod_model)
 
 style_files = os.listdir('./style_images')
 style_options = {
@@ -52,26 +48,19 @@ style_options = {
 lrs = np.linspace(0.015, 0.075, 3).tolist()
 img_size = 512
 
-cached_style_features = {}
-for style_name, style_img_path in style_options.items():
-    style_img = preprocess_img(style_img_path, img_size)[0].to(device)
-    with torch.no_grad():
-        style_features = model(style_img)
-    cached_style_features[style_name] = style_features 
+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()
+}
 
-@spaces.GPU(duration=30)
-def run(content_image, style_name, style_strength=len(lrs), optim_name='AdamW', apply_to_background=False):
+@spaces.GPU(duration=15)
+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]
     
-    if optim_name == 'AdamW': 
-        optim_caller = optim.AdamW
-    elif optim_name == 'L-BFGS': 
-        optim_caller = optim.LBFGS
-    
     print('-'*30)
     print(datetime.now(timezone.utc) - timedelta(hours=5)) # EST
     
@@ -84,16 +73,11 @@ def run(content_image, style_name, style_strength=len(lrs), optim_name='AdamW',
         style_features=style_features,
         lr=lrs[style_strength-1],
         apply_to_background=apply_to_background,
-        optim_caller=optim_caller,
     )
-    et = time.time()
-    print(f'{et-st:.2f}s')
+    print(f'{time.time()-st:.2f}s')
     
     yield postprocess_img(generated_img, original_size)
 
-def set_slider(value):
-    return gr.update(value=value)
-
 css = """
 #container {
     margin: 0 auto;
@@ -111,13 +95,10 @@ with gr.Blocks(css=css) as demo:
                 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)
-            with gr.Accordion(label='Advanced Options', open=False):
-                optim_dropdown = gr.Radio(choices=['AdamW', 'L-BFGS'], label='Optimizer', value='AdamW', type='value')
             submit_button = gr.Button('Submit', variant='primary')
             
             examples = gr.Examples(
                 examples=[
-                    ['./content_images/Surfer.jpg', 'Starry Night'],
                     ['./content_images/GoldenRetriever.jpg', 'Great Wave'],
                     ['./content_images/CameraGirl.jpg', 'Bokeh']
                 ],
@@ -140,7 +121,7 @@ with gr.Blocks(css=css) as demo:
         
     submit_button.click(
         fn=run, 
-        inputs=[content_image, style_dropdown, style_strength_slider, optim_dropdown, apply_to_background_checkbox], 
+        inputs=[content_image, style_dropdown, style_strength_slider, apply_to_background_checkbox], 
         outputs=output_image
     ).then(
         fn=save_image,

utils.py

@@ -2,6 +2,7 @@ from PIL import Image
 
 import torch
 import torchvision.transforms as transforms
+from safetensors.torch import load_file
 
 def preprocess_img(img, img_size, normalize=False):
     if type(img) == str: img = Image.open(img)
@@ -33,4 +34,11 @@ def postprocess_img(img, original_size, normalize=False):
     
     img = transforms.ToPILImage()(img)
     img = img.resize(original_size, Image.Resampling.LANCZOS)
-    return img
+    return img
+
+def load_model_without_module(model, model_path, device):
+    state_dict = {
+        k[7:] if k.startswith('module.') else k: v 
+        for k, v in load_file(model_path, device=device).items()
+    }
+    model.load_state_dict(state_dict)