Update app.py

app.py

@@ -8,14 +8,18 @@ import numpy as np
 import requests
 from io import BytesIO
 
-def load_image(img_path, max_size=400, shape=None):
-    ''' Load in and transform an image, making sure the image
-       is <= 400 pixels in the x-y dims.'''
-    if "http" in img_path:
-        response = requests.get(img_path)
-        image = Image.open(BytesIO(response.content)).convert('RGB')
+def load_image(img_input, max_size=400, shape=None):
+    if isinstance(img_input, np.ndarray):
+        # Convert numpy array to PIL Image
+        image = Image.fromarray(img_input.astype('uint8'), 'RGB')
+    elif isinstance(img_input, str):
+        if "http" in img_input:
+            response = requests.get(img_input)
+            image = Image.open(BytesIO(response.content)).convert('RGB')
+        else:
+            image = Image.open(img_input).convert('RGB')
     else:
-        image = Image.open(img_path).convert('RGB')
+        raise ValueError("Unsupported input type. Expected numpy array or string.")
     
     # large images will slow down processing
     if max(image.size) > max_size:
@@ -97,16 +101,6 @@ def gram_matrix(tensor):
     
     return gram
 
-#Load VGG19 model
-vgg = models.vgg19(pretrained=True).features
-for param in vgg.parameters():
-    param.requires_grad_(False)
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-vgg.to(device)
-
-# Helper functions (load_image, im_convert, get_features, gram_matrix)
-# ... (Include the helper functions you provided earlier here)
-
 def style_transfer(content_image, style_image, alpha, beta, conv1_1, conv2_1, conv3_1, conv4_1, conv5_1, steps):
     content = load_image(content_image).to(device)
     style = load_image(style_image, shape=content.shape[-2:]).to(device)
@@ -161,6 +155,14 @@ examples = [
      "https://huggingface.co/spaces/muneebable/vgg-style-transfer/resolve/main/assets/style_3.jpg"],
 ]
 
+#Load VGG19 model
+vgg = models.vgg19(pretrained=True).features
+for param in vgg.parameters():
+    param.requires_grad_(False)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+vgg.to(device)
+
+
 # Gradio interface
 with gr.Blocks() as demo:
     gr.Markdown("# Neural Style Transfer")