Update app.py

app.py

@@ -8,12 +8,15 @@ import matplotlib.pyplot as plt
 from pathlib import Path
 
 # ================================================
-# REQUIRED STUB CLASSES / FUNCTIONS FOR LOADING
+# STUBS FOR CUSTOM CLASSES / FUNCTIONS
+# These match all custom classes used in the model
 # ================================================
 
-def _inner(*args, **kwargs):
-    return None
+# Functions
+def _inner(*args, **kwargs): return None
+def style_loss(*args, **kwargs): return None
 
+# Layers
 class ReflectionLayer(nn.Module):
     def __init__(self): super().__init__()
     def forward(self, x): return x
@@ -49,10 +52,10 @@ learner = from_pretrained_fastai(REPO_ID)
 # ================================================
 
 def infer(img):
-    # Convert to PILImage if needed
+    # Ensure PILImage
     if not isinstance(img, PILImage):
         img = PILImage.create(img)
-    
+
     pred = learner.predict(img)
     image = pred[0].numpy()
     image = image.transpose((1, 2, 0))  # CHW -> HWC