Upload folder using huggingface_hub

app.py

@@ -4,6 +4,7 @@ from datetime import datetime, timezone, timedelta
 
 import spaces
 import torch
+import torchvision.models as models
 import numpy as np
 import gradio as gr
 from gradio_imageslider import ImageSlider
@@ -21,6 +22,9 @@ if device == 'cuda': print('CUDA DEVICE:', torch.cuda.get_device_name())
 model = VGG_19().to(device).eval()
 for param in model.parameters():
     param.requires_grad = False
+segmentation_model = models.segmentation.deeplabv3_resnet101(
+    weights='DEFAULT'
+).to(device).eval()
 
 style_files = os.listdir('./style_images')
 style_options = {' '.join(style_file.split('.')[0].split('_')): f'./style_images/{style_file}' for style_file in style_files}
@@ -51,6 +55,7 @@ def run(content_image, style_name, style_strength=5, apply_to_background=False,
     st = time.time()
     generated_img = inference(
         model=model,
+        segmentation_model=segmentation_model,
         content_image=content_img,
         style_features=style_features,
         lr=lrs[style_strength-1],

inference.py

@@ -28,7 +28,6 @@ def _compute_loss(generated_features, content_features, style_features, resized_
         else:
             G = _gram_matrix(gf)
             A = _gram_matrix(sf)
-            style_loss += w_l * F.mse_loss(G, A)
         style_loss += w_l * F.mse_loss(G, A)
         
     return alpha * content_loss + beta * style_loss
@@ -36,6 +35,7 @@ def _compute_loss(generated_features, content_features, style_features, resized_
 def inference(
     *,
     model,
+    segmentation_model,
     content_image,
     style_features,
     apply_to_background,
@@ -53,10 +53,7 @@ def inference(
         content_features = model(content_image)
 
         resized_bg_masks = []        
-        if apply_to_background:
-            segmentation_model = models.segmentation.deeplabv3_resnet101(weights='DEFAULT').eval()
-            segmentation_model = segmentation_model.to(content_image.device)
-            
+        if apply_to_background:            
             segmentation_output = segmentation_model(content_image)['out']
             segmentation_mask = segmentation_output.argmax(dim=1)