Update app.py

app.py

@@ -5,21 +5,24 @@ from torchvision import transforms
 from PIL import Image
 from src.lightning_module import StyleTransferModule
 
-
 MODEL_URL = "https://huggingface.co/Michal-Raszkowski/adain-style-transfer/resolve/main/style-transfer-best-v2.ckpt?download=true"
 CHECKPOINT_PATH = "model.ckpt"
 
 def download_model_if_missing():
     if not os.path.exists(CHECKPOINT_PATH):
+        print("Downloading model checkpoint...")
         torch.hub.download_url_to_file(MODEL_URL, CHECKPOINT_PATH)
 
 def load_model():
     download_model_if_missing()
-    model = StyleTransferModule.load_from_checkpoint(CHECKPOINT_PATH, map_location="cpu")
+    # Loading to CPU by default; change to "cuda" if GPU is available
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = StyleTransferModule.load_from_checkpoint(CHECKPOINT_PATH, map_location=device)
     model.eval()
-    return model
+    return model, device
 
-model = load_model()
+# Initialize model and get target device
+model, device = load_model()
 
 def stylize(content_image, style_image, alpha):
     if content_image is None or style_image is None:
@@ -30,34 +33,46 @@ def stylize(content_image, style_image, alpha):
         transforms.ToTensor()
     ])
     
-    c = transform(content_image).unsqueeze(0)
-    s = transform(style_image).unsqueeze(0)
+    # Transform and push tensors to the correct device (CPU/GPU)
+    c = transform(content_image).unsqueeze(0).to(device)
+    s = transform(style_image).unsqueeze(0).to(device)
     
     with torch.no_grad():
         generated_tensor, _ = model(c, s, alpha=alpha)
     
-    generated_tensor = torch.clamp(generated_tensor, 0, 1)
-    result_image = transforms.ToPILImage()(generated_tensor.squeeze(0))
+    # Bring back to CPU, clamp values, remove batch dimension, and convert to PIL
+    generated_tensor = torch.clamp(generated_tensor, 0, 1).cpu().squeeze(0)
+    result_image = transforms.ToPILImage()(generated_tensor)
     
     return result_image
 
+# Build the Gradio Interface
 with gr.Blocks(title="Style Transfer Demo", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("Neural Style Transfer")
-    gr.Markdown("Upload content and style images to combine them.")
+    gr.Markdown("# Neural Style Transfer")
+    gr.Markdown("Upload a content image and a style image to combine them using AdaIN.")
     
     with gr.Row():
         with gr.Column():
-            input_content = gr.Image(label="Content image", type="pil", height=300)
-            input_style = gr.Image(label="Style image", type="pil", height=300)
-            slider = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, label="Style strenght.")
+            input_content = gr.Image(label="Content Image", type="pil", height=300)
+            input_style = gr.Image(label="Style Image", type="pil", height=300)
+            slider = gr.Slider(minimum=0.0, maximum=1.0, value=1.0, step=0.1, label="Style Strength")
             btn = gr.Button("Generate", variant="primary")
         
         with gr.Column():
-            output = gr.Image(label="Output", type="pil")
+            output = gr.Image(label="Output Image", type="pil")
     
-    btn.click(fn=stylize, inputs=[input_content, input_style, slider], outputs=output)
+    # Set up the click event listener
+    btn.click(
+        fn=stylize, 
+        inputs=[input_content, input_style, slider], 
+        outputs=output
+    )
     
-    #gr.Examples(examples=[["examples/c.jpg", "examples/s.jpg", 1.0]], inputs=[input_content, input_style, slider])
+    # Optional: Uncomment if you want to include default examples
+    # gr.Examples(
+    #     examples=[["examples/c.jpg", "examples/s.jpg", 1.0]], 
+    #     inputs=[input_content, input_style, slider]
+    # )
 
 if __name__ == "__main__":
     demo.launch()