Update app.py

app.py

@@ -0,0 +1,120 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import models, transforms
+from torchvision.models import VGG19_Weights
+from PIL import Image
+import gradio as gr
+
+# ✅ Use GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print("Using device:", device)
+
+# --- Image Utilities ---
+def load_image(img, max_size=512):
+    transform = transforms.Compose([
+        transforms.Resize(max_size),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                             std=[0.229, 0.224, 0.225])
+    ])
+    image = img.convert('RGB')
+    image = transform(image).unsqueeze(0)
+    return image.to(device)
+
+def tensor_to_image(tensor):
+    unnormalize = transforms.Normalize(
+        mean=[-0.485 / 0.229, -0.456 / 0.224, -0.406 / 0.225],
+        std=[1 / 0.229, 1 / 0.224, 1 / 0.225]
+    )
+    image = tensor.clone().detach().squeeze(0)
+    image = unnormalize(image)
+    image = torch.clamp(image, 0, 1)
+    return transforms.ToPILImage()(image)
+
+# --- Style Transfer Utilities ---
+def gram_matrix(tensor):
+    b, c, h, w = tensor.size()
+    features = tensor.view(b * c, h * w)
+    return torch.mm(features, features.t())
+
+class StyleTransferNet(nn.Module):
+    def __init__(self, style_img, content_img):
+        super().__init__()
+        weights = VGG19_Weights.DEFAULT
+        self.vgg = models.vgg19(weights=weights).features.to(device).eval()
+        self.style_img = style_img
+        self.content_img = content_img
+        self.content_layers = ['conv_4']
+        self.style_layers = ['conv_1', 'conv_3', 'conv_5', 'conv_9']
+
+    def get_features(self, x):
+        features = {}
+        i = 0
+        for layer in self.vgg.children():
+            x = layer(x)
+            if isinstance(layer, nn.Conv2d):
+                i += 1
+                name = f'conv_{i}'
+                if name in self.content_layers + self.style_layers:
+                    features[name] = x
+        return features
+
+    def forward(self, input_img, steps=300, style_weight=1e6, content_weight=0.25):
+        input_img = input_img.clone().requires_grad_(True)
+        optimizer = optim.LBFGS([input_img])
+
+        style_features = self.get_features(self.style_img)
+        content_features = self.get_features(self.content_img)
+        style_grams = {k: gram_matrix(v) for k, v in style_features.items()}
+
+        run = [0]
+        while run[0] <= steps:
+            def closure():
+                optimizer.zero_grad()
+                target_features = self.get_features(input_img)
+                style_loss = 0
+                content_loss = 0
+
+                for layer in self.style_layers:
+                    target_feature = target_features[layer]
+                    target_gram = gram_matrix(target_feature)
+                    style_gram = style_grams[layer]
+                    style_loss += torch.mean((target_gram - style_gram)**2)
+
+                for layer in self.content_layers:
+                    target_feature = target_features[layer]
+                    content_feature = content_features[layer]
+                    content_loss += torch.mean((target_feature - content_feature)**2)
+
+                total_loss = style_weight * style_loss + content_weight * content_loss
+                total_loss.backward(retain_graph=True)
+                run[0] += 1
+                return total_loss
+
+            optimizer.step(closure)
+        return input_img
+
+# --- Gradio App ---
+def style_transfer_app(content_img, style_img, content_weight, style_weight, steps):
+    content = load_image(content_img)
+    style = load_image(style_img)
+    model = StyleTransferNet(style, content)
+    output = model(content, steps=int(steps), content_weight=content_weight, style_weight=style_weight)
+    return tensor_to_image(output)
+
+# --- Launch Interface ---
+gr.Interface(
+    fn=style_transfer_app,
+    inputs=[
+        gr.Image(type="pil", label="🖼️ Content Image"),
+        gr.Image(type="pil", label="🎨 Style Image"),
+        gr.Slider(0.05, 1.0, value=0.25, step=0.05, label="Content Weight"),
+        gr.Slider(1e5, 5e6, value=1e6, step=1e5, label="Style Weight"),
+        gr.Slider(50, 500, value=300, step=50, label="Steps")
+    ],
+    outputs=gr.Image(type="pil", label="🧠 Stylized Output"),
+    title="🧠 AI Neural Style Transfer Lab",
+    description="Upload a content image and a style image. Then tweak the controls below to explore the balance between structure and stylization. Powered by PyTorch + VGG19.",
+    allow_flagging="never"
+).launch(share=True)