Update main

app.py

@@ -1,3 +1,4 @@
+
 import torch
 import torch.nn as nn
 import torch.optim as optim
@@ -5,13 +6,14 @@ from torchvision import models, transforms
 from torchvision.models import VGG19_Weights
 from PIL import Image
 import gradio as gr
+import time
 
 # ✅ 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):
+def load_image(img, max_size=384):
     transform = transforms.Compose([
         transforms.Resize(max_size),
         transforms.ToTensor(),
@@ -60,48 +62,55 @@ class StyleTransferNet(nn.Module):
                     features[name] = x
         return features
 
-    def forward(self, input_img, steps=300, style_weight=1e6, content_weight=0.25):
+    def forward(self, input_img, steps=100, style_weight=1e6, content_weight=1e5):
         input_img = input_img.clone().requires_grad_(True)
-        optimizer = optim.LBFGS([input_img])
+        optimizer = optim.Adam([input_img], lr=0.02)
 
         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)
+        for step in range(steps):
+            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()
+            optimizer.step()
+
         return input_img
 
 # --- Gradio App ---
-def style_transfer_app(content_img, style_img, content_weight, style_weight, steps):
+def style_transfer_app(content_img, style_img, content_weight_ui, style_weight_ui, steps):
+    start_time = time.time()
     content = load_image(content_img)
     style = load_image(style_img)
+
+    # Map intuitive UI weights (1-10) to actual values
+    content_weight = content_weight_ui * 1e5
+    style_weight = style_weight_ui * 1e6
+
     model = StyleTransferNet(style, content)
     output = model(content, steps=int(steps), content_weight=content_weight, style_weight=style_weight)
-    return tensor_to_image(output)
+    stylized = tensor_to_image(output)
+    elapsed = round(time.time() - start_time)
+
+    # Estimated time display
+    estimate_note = f"🕒 Estimated processing time: {elapsed} seconds for {steps} steps."
+    return stylized, estimate_note
 
 # --- Launch Interface ---
 gr.Interface(
@@ -109,14 +118,15 @@ gr.Interface(
     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")
+        gr.Slider(1, 10, value=1, step=1, label="Content Weight (1 = weak structure, 10 = strong)"),
+        gr.Slider(1, 10, value=6, step=1, label="Style Weight (1 = subtle, 10 = strong style)"),
+        gr.Slider(50, 300, value=100, step=50, label="Steps (speed vs quality)")
     ],
-    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.",
+    outputs=[
+        gr.Image(type="pil", label="🧠 Stylized Output"),
+        gr.Textbox(label="⏱️ Time Info")
+    ],
+    title="🎨 Fast AI Neural Style Transfer",
+    description="Upload content and style images, then tune how much structure vs style you want. Powered by PyTorch + VGG19.",
     allow_flagging="never"
 ).launch(share=True)
-
-gr.Interface(...).launch()