Update app.py

app.py

@@ -1,87 +1,71 @@
 import torch
 import torch.nn as nn
+import torchvision.transforms as transforms
+from PIL import Image
 import gradio as gr
-import numpy as np
-from huggingface_hub.fastai_utils import from_pretrained_fastai
-from fastai.vision.all import *
-from pathlib import Path
+from torchvision.utils import save_image
 
 # ================================================
-# STUBS FOR CUSTOM CLASSES / FUNCTIONS
+# FAST NEURAL STYLE MODEL (AdaIN) - Pytorch Hub
 # ================================================
 
-# Functions
-def _inner(*args, **kwargs): return None
-def style_loss(*args, **kwargs): return None
-def act_loss(*args, **kwargs): return None
-
-# Layers
-class ReflectionLayer(nn.Module):
-    def __init__(self): super().__init__()
-    def forward(self, x): return x
-
-class ConvLayer(nn.Module):
-    def __init__(self): super().__init__()
-    def forward(self, x): return x
-
-class ResidualBlock(nn.Module):
-    def __init__(self): super().__init__()
-    def forward(self, x): return x
-
-class UpsampleConvLayer(nn.Module):
-    def __init__(self): super().__init__()
-    def forward(self, x): return x
-
-class TransformerNet(nn.Module):
-    def __init__(self): super().__init__()
-    def forward(self, x): return x
-
-class FeatureLoss:
-    def __init__(self, *args, **kwargs): pass
+# Tutaj pobieramy model z PyTorch Hub (dynamic style transfer)
+# Uwaga: model pobiera content + style image
+model = torch.hub.load('pytorch/examples', 'fast_neural_style', source='github', model='candy')  # przykładowy styl
+model.eval()
 
 # ================================================
-# LOAD FASTAI MODEL
+# UTILS
 # ================================================
 
-REPO_ID = "hugginglearners/fastai-style-transfer"
-learner = from_pretrained_fastai(REPO_ID)
+def preprocess(img, size=512):
+    """Konwersja PIL -> Tensor"""
+    transform = transforms.Compose([
+        transforms.Resize(size),
+        transforms.ToTensor(),
+        transforms.Lambda(lambda x: x.mul(255))
+    ])
+    img_t = transform(img).unsqueeze(0)  # dodaj batch dim
+    return img_t
+
+def postprocess(tensor):
+    """Tensor -> PIL Image"""
+    tensor = tensor.clamp(0, 255).squeeze(0)
+    tensor = tensor / 255
+    return transforms.ToPILImage()(tensor)
 
 # ================================================
-# INFERENCE FUNCTION
+# INFERENCE
 # ================================================
 
-def infer(img):
-    # Convert to PILImage if needed
-    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
-    image = np.clip(image*255, 0, 255).astype(np.uint8)
-    return image
+def infer(content_img, style_img):
+    # Preprocess
+    content = preprocess(content_img)
+    style = preprocess(style_img)
+    
+    # Użyj modelu dynamic style transfer
+    with torch.no_grad():
+        output = model(content, style)  # content + style
+    
+    return postprocess(output)
 
 # ================================================
 # GRADIO INTERFACE
 # ================================================
 
-EXAMPLES_PATH = Path("./examples")
-examples = [str(EXAMPLES_PATH / f.name) for f in EXAMPLES_PATH.iterdir()]
-
-title = "🎨 FastAI Style Transfer"
-description = "Transform images using a neural style transfer model."
-article = 'Author: <a href="https://huggingface.co/geninhu">Nhu Hoang</a>.'
+title = "🎨 Dynamic Neural Style Transfer"
+description = "Upload a content image and a style image to apply style transfer dynamically."
 
 app = gr.Interface(
     fn=infer,
-    inputs=gr.Image(type="pil", label="Upload Image"),
+    inputs=[
+        gr.Image(type="pil", label="Content Image"),
+        gr.Image(type="pil", label="Style Image")
+    ],
     outputs=gr.Image(label="Stylized Image"),
-    examples=examples,
     title=title,
     description=description,
-    article=article,
-    allow_flagging="never",
-    cache_examples=False
+    allow_flagging="never"
 )
 
 if __name__ == "__main__":