gradio-fix (#2)

- gradio update (ff63d1bb8456f8b2612f9096bc239f48d3220790)

.gitignore

@@ -0,0 +1,3 @@
+.venv/
+.gradio/
+__pycache__/

app.py

@@ -4,39 +4,22 @@ import clip
 from torchvision import transforms
 from style_net import StyleNet
 from dataset import FlickrStreamer
-from util import DEVICE
-from clip_styler import ClipStyler
+from clip_styler import style_transfer
 
 # --- CONFIGURATION ---
-#DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 INDEX_PATH = "flickr_embeddings.pt"
 SEARCH_LIMIT = 5000  # Must match the indexer limit!
 # ---------------------
 
-UI_EPOCHS = 100
-STYLE_PRESETS = {
-    "Sketch": {
-        "prompt": "charcoal sketch style",
-        "source": "a Photo",
-    },
-    "Van Gogh": {
-        "prompt": "Van Gogh painting style",
-        "source": "a Photo",
-    },
-    "Cyberpunk": {
-        "prompt": "neon cyberpunk style",
-        "source": "a Photo",
-    },
-}
-
 print(f"System starting on {DEVICE}...")
 
 # 1. LOAD MODELS
-# A. Retrieval Model (CLIP)
-clip_model, _ = clip.load("ViT-B/32", device=DEVICE)
+# A. Editing Model
+#style_net = StyleNet().to(DEVICE)
 
-# B. Editing Model (CLIP-guided optimizer)
-clip_styler = ClipStyler(device=DEVICE, num_epochs=UI_EPOCHS)
+# B. Retrieval Model (CLIP)
+#clip_model, preprocess = clip.load("ViT-B/32", device=DEVICE)
 
 # 2. LOAD SEARCH INDEX
 print("Loading Search Index...")
@@ -78,25 +61,39 @@ def search_flickr(text_query):
             
         return results
 
-def edit_image(image, style_choice):
-    if image is None or style_choice is None:
-        return None
-
-    preset = STYLE_PRESETS.get(style_choice)
-    if preset is None:
-        return image
+# --- FUNCTION: EDIT ---
+tf_edit = transforms.Compose([
+    transforms.Resize((256, 256)),
+    transforms.ToTensor()
+])
 
+def edit_image(image, style_choice, num_steps):
+    if image is None: return None
+    
+    weights_map = {
+        "Sketch": "model_sketch.pth",
+        "Van Gogh": "model_painting.pth",
+        "Cyberpunk": "model_cyberpunk.pth"
+    }
+    detail_prompt = {
+        "Sketch": "charcoal sketch",
+        "Van Gogh": "Van Gogh painting",
+        "Cyberpunk": "neon cyberpunk",
+    }
+    '''
     try:
-        return clip_styler.stylize(
-            image=image,
-            prompt=preset["prompt"],
-            source=preset.get("source", "a Photo"),
-            num_epochs=UI_EPOCHS,
-            log_interval=0,
-        )
-    except Exception as exc:
-        print(f"[edit_image] Failed to stylize image: {exc}")
+        style_net.load_state_dict(torch.load(weights_map.get(style_choice), map_location=DEVICE))
+    except Exception:
         return image 
+    
+    input_tensor = tf_edit(image).unsqueeze(0).to(DEVICE)
+    with torch.no_grad():
+        output_tensor = style_net(input_tensor)
+        
+    output_img = transforms.ToPILImage()(output_tensor.squeeze().cpu())
+    return output_img
+    '''
+    return style_transfer(image, detail_prompt[style_choice], num_steps=num_steps)
 
 # --- THE UI ---
 with gr.Blocks() as demo:
@@ -118,11 +115,18 @@ with gr.Blocks() as demo:
         with gr.Row():
             im_in = gr.Image(type="pil", label="Input Image")
             style = gr.Dropdown(["Sketch", "Van Gogh", "Cyberpunk"], label="Choose Style")
+            steps_slider = gr.Slider(
+                minimum=50,
+                maximum=200,
+                value=100,
+                step=10,
+                label="Number of Steps (more = better quality but slower)"
+            )
         
         btn_edit = gr.Button("Apply")
         im_out = gr.Image(label="Result")
         
-        btn_edit.click(edit_image, inputs=[im_in, style], outputs=im_out)
+        btn_edit.click(edit_image, inputs=[im_in, style, steps_slider], outputs=im_out)
 
 if __name__ == "__main__":
     demo.launch(share=True)

clip_styler.py

@@ -1,212 +1,166 @@
-from __future__ import annotations
-
-from typing import Optional
-
-from PIL import Image
 import torch
-import torch.optim as optim
-from torchvision import transforms, models
-from torchvision.transforms.functional import adjust_contrast
+from PIL import Image
+import torchvision.transforms as transforms
+from util import *
+from cnn1 import UNet
 from transformers import CLIPModel, AutoTokenizer
+from torchvision.transforms.functional import adjust_contrast
 
-from cnn1 import UNet
-from util import (
-    DEVICE,
-    IMG_SIZE,
-    LR,
-    CROP_SIZE,
-    RESIZE,
-    NUM_EPOCHS,
-    NUM_CROPS,
-    PATCH_THRESHOLD,
-    L_TV,
-    L_PATCH,
-    L_DIR,
-    L_CONTENT,
-    normalize,
-    clip_normalize,
-    get_features,
-    prompt_ensemble,
-    get_image_prior_losses,
-)
-
-
-class ClipStyler:
+clip_model = None
+tokenizer = None
+vgg19 = None
+cnn_model = None
+
+def load_models():
+    """Load models once and cache them"""
+    global clip_model, tokenizer, vgg19, cnn_model
+    
+    if clip_model is None:
+        clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(DEVICE)
+    
+    if tokenizer is None:
+        tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+    
+    if vgg19 is None:
+        vgg19 = models.vgg19(pretrained=True).features.to(DEVICE)
+        for param in vgg19.parameters():
+            param.requires_grad_(False)
+    
+    if cnn_model is None:
+        cnn_model = UNet(text_dim=512).to(DEVICE)
+    
+    return clip_model, tokenizer, vgg19, cnn_model
+
+def reset_cnn_model():
+    """Reset CNN model weights for each new image"""
+    global cnn_model
+    if cnn_model is not None:
+        # Reinitialize the model
+        for layer in cnn_model.modules():
+            if isinstance(layer, (torch.nn.Conv2d, torch.nn.Linear)):
+                torch.nn.init.xavier_uniform_(layer.weight)
+                if layer.bias is not None:
+                    torch.nn.init.zeros_(layer.bias)
+
+def style_transfer(img, prompt, num_steps=NUM_EPOCHS, source="a Photo"):
     """
-    Wraps the CLIP-guided optimization loop from main.py so it can be reused
-    by both the CLI entrypoint and the Gradio application.
+    Apply style transfer to an uploaded image
+    
+    Args:
+        image: PIL Image uploaded by user
+        prompt: Text style prompt
+        num_steps: Number of optimization steps (fewer = faster, more = better quality)
+    
+    Returns:
+        Stylized PIL Image
     """
-
-    def __init__(
-        self,
-        device: Optional[torch.device] = None,
-        num_epochs: int = NUM_EPOCHS,
-    ) -> None:
-        self.device = device or DEVICE
-        self.num_epochs = num_epochs
-
-        # Text/image encoders
-        self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(
-            self.device
-        )
-        self.clip_model.eval()
-        for param in self.clip_model.parameters():
-            param.requires_grad_(False)
-
-        self.tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")
-
-        # Perceptual backbone
-        self.vgg19 = models.vgg19(pretrained=True).features.to(self.device).eval()
-        for param in self.vgg19.parameters():
-            param.requires_grad_(False)
-
-        # Image helpers
-        self.image_transform = transforms.Compose(
-            [transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.ToTensor()]
-        )
-        self.to_pil = transforms.ToPILImage()
-        self.crop_transform = transforms.RandomCrop(CROP_SIZE)
-        self.augment_transform = transforms.Compose(
-            [
-                transforms.RandomPerspective(fill=0, p=1, distortion_scale=0.5),
-                transforms.Resize(RESIZE),
-            ]
-        )
-
-    def stylize(
-        self,
-        image: Image.Image,
-        prompt: str,
-        source: str = "a Photo",
-        num_epochs: Optional[int] = None,
-        log_interval: int = 20,
-    ) -> Image.Image:
-        """
-        Run the CLIP-guided optimization loop and return the stylized PIL image.
-        """
-        if image is None:
-            raise ValueError("Image must be provided for stylization.")
-
-        epochs = num_epochs if num_epochs is not None else self.num_epochs
-        base_img = self._pil_to_tensor(image)
-        content_features = get_features(normalize(base_img), self.vgg19)
-
-        cnn_model = UNet(text_dim=512).to(self.device)
-        cnn_model.train()
-
-        optimizer = optim.Adam(cnn_model.parameters(), lr=LR)
-        scheduler = torch.optim.lr_scheduler.StepLR(
-            optimizer, step_size=100, gamma=0.5
-        )
-
-        with torch.no_grad():
-            text_features = self._encode_text(prompt)
-            source_features = self._encode_text(source)
-            base_img_features = self._image_features(base_img)
-
-        style_direction = text_features - source_features
-        style_direction /= style_direction.norm(dim=-1, keepdim=True)
-
-        final_target = base_img
-        for epoch in range(epochs + 1):
-            scheduler.step()
-            target = cnn_model(base_img, text_embedding=text_features)
-            final_target = target
-
-            target_features = get_features(normalize(target), self.vgg19)
-            content_loss = torch.mean(
-                (target_features["conv4_2"] - content_features["conv4_2"]) ** 2
-            )
-            content_loss += torch.mean(
-                (target_features["conv5_2"] - content_features["conv5_2"]) ** 2
-            )
-
-            img_aug = self._augment_crops(target)
-            img_aug_features = self.clip_model.get_image_features(
-                pixel_values=clip_normalize(img_aug)
-            )
-            img_aug_features /= img_aug_features.norm(dim=-1, keepdim=True)
-
-            img_direction = img_aug_features - base_img_features
-            img_direction /= img_direction.norm(dim=-1, keepdim=True)
-
-            patch_text_direction = style_direction.repeat(img_direction.size(0), 1)
-            patch_text_direction /= patch_text_direction.norm(dim=-1, keepdim=True)
-
-            loss_patch = 1 - torch.cosine_similarity(
-                img_direction, patch_text_direction, dim=1
-            )
-            loss_patch = torch.where(
-                loss_patch < PATCH_THRESHOLD, torch.zeros_like(loss_patch), loss_patch
-            ).mean()
-
-            glob_features = self.clip_model.get_image_features(
-                pixel_values=clip_normalize(target)
-            )
-            glob_features /= glob_features.norm(dim=-1, keepdim=True)
-
-            glob_direction = glob_features - base_img_features
-            glob_direction /= glob_direction.norm(dim=-1, keepdim=True)
-
-            loss_glob = (1 - torch.cosine_similarity(glob_direction, style_direction, dim=1)).mean()
-
-            loss_tv = L_TV * get_image_prior_losses(target)
-
-            total_loss = (
-                (L_PATCH * loss_patch)
-                + (L_DIR * loss_glob)
-                + (L_CONTENT * content_loss)
-                + loss_tv
-            )
-
-            optimizer.zero_grad()
-            total_loss.backward()
-            optimizer.step()
-
-            '''
-            if log_interval and epoch % log_interval == 0:
-                print(
-                    f"[ClipStyler] Epoch {epoch}: total={total_loss.item():.4f}, "
-                    f"content={content_loss.item():.4f}, patch={loss_patch.item():.4f}, "
-                    f"dir={loss_glob.item():.4f}, tv={loss_tv.item():.4f}"
-                )
-            '''
-
-        return self._tensor_to_pil(final_target)
-
-    def _pil_to_tensor(self, image: Image.Image) -> torch.Tensor:
-        if image.mode != "RGB":
-            image = image.convert("RGB")
-        tensor = self.image_transform(image).unsqueeze(0)
-        return tensor.to(self.device)
-
-    def _tensor_to_pil(self, tensor: torch.Tensor) -> Image.Image:
-        tensor = tensor.detach().clone().cpu().clamp(0, 1).squeeze(0)
-        tensor = adjust_contrast(tensor, 1.5)
-        return self.to_pil(tensor)
-
-    def _augment_crops(self, target: torch.Tensor) -> torch.Tensor:
-        crops = []
-        for _ in range(NUM_CROPS):
-            crop = self.crop_transform(target)
-            crop = self.augment_transform(crop)
-            crops.append(crop)
-        return torch.cat(crops, dim=0)
-
-    @torch.no_grad()
-    def _encode_text(self, text: str) -> torch.Tensor:
-        prompts = prompt_ensemble(text)
-        tokens = self.tokenizer(prompts, padding=True, return_tensors="pt").to(
-            self.device
-        )
-        text_features = self.clip_model.get_text_features(**tokens)
-        text_features = text_features.mean(dim=0, keepdim=True)
+    if img is None:
+        return None
+    
+    img = load_image(img).to(DEVICE)
+    
+    # Load models
+    clip_model, tokenizer, vgg19, cnn_model = load_models()
+    
+    # Reset CNN model for fresh start on each image
+    reset_cnn_model()
+    
+    # image features
+    features = get_features(normalize(img), vgg19)
+
+    # cnn model
+    #cnn_model = UNet().to(DEVICE)
+
+    # load ADAM
+    optimizer = optim.Adam(cnn_model.parameters(), lr=LR)
+    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=100, gamma=0.5)
+
+    # data augmentation
+    crop = transforms.Compose([transforms.RandomCrop(CROP_SIZE)])
+    augment = transforms.Compose([transforms.RandomPerspective(fill=0, p=1, distortion_scale=0.5), transforms.Resize(RESIZE)])
+
+    # initialize variables
+    content_loss_epoch = []
+    style_loss_epoch = []
+    total_loss_epoch = []
+    output_image = img
+    mean_img = torch.mean(img, dim=(2, 3), keepdim=False).squeeze(0)
+    mean_img = [mean_img[0].item(), mean_img[1].item(), mean_img[2].item()]
+    target = img.clone().requires_grad_(True).to(DEVICE)
+
+    with torch.no_grad():
+        edited_text = prompt_ensemble(prompt)
+        text_tokens = tokenizer(edited_text, padding=True, return_tensors="pt").to(DEVICE)
+        text_features = clip_model.get_text_features(**text_tokens)
+        text_features = text_features.mean(axis=0, keepdim=True)
         text_features /= text_features.norm(dim=-1, keepdim=True)
-        return text_features
-
-    @torch.no_grad()
-    def _image_features(self, tensor: torch.Tensor) -> torch.Tensor:
-        feats = self.clip_model.get_image_features(pixel_values=clip_normalize(tensor))
-        feats /= feats.norm(dim=-1, keepdim=True)
-        return feats
 
+        edited_source = prompt_ensemble(source)
+        source_tokens = tokenizer(edited_source, padding=True, return_tensors="pt").to(DEVICE)
+        source_features = clip_model.get_text_features(**source_tokens)
+        source_features = source_features.mean(axis=0, keepdim=True)
+        source_features /= source_features.norm(dim=-1, keepdim=True)
+
+        img_features = clip_model.get_image_features(pixel_values=clip_normalize(img))
+        img_features /= (img_features.clone().norm(dim=-1, keepdim=True))
+
+    for epoch in range(num_steps + 1):
+        scheduler.step()
+        #target = cnn_model(img).requires_grad_(True).to(DEVICE)
+        target = cnn_model(img, text_embedding=text_features).to(DEVICE)
+
+        target_features = get_features(normalize(target), vgg19)
+
+        content_loss = 0.0
+        content_loss += torch.mean((target_features['conv4_2'] - features['conv4_2']) ** 2)
+        content_loss += torch.mean((target_features['conv5_2'] - features['conv5_2']) ** 2)
+
+        loss_patch = 0 
+        img_proc = []
+        for n in range(NUM_CROPS):
+            target_crop = crop(target)
+            target_crop = augment(target_crop)
+            img_proc.append(target_crop)
+
+        img_proc = torch.cat(img_proc, dim=0)
+        img_aug = img_proc
+
+        img_aug_features = clip_model.get_image_features(pixel_values=clip_normalize(img_aug))
+        img_aug_features /= (img_aug_features.clone().norm(dim=-1, keepdim=True))
+
+        img_direction = img_aug_features - img_features
+        img_direction /= img_direction.clone().norm(dim=-1, keepdim=True)
+
+        text_direction = (text_features - source_features).repeat(img_aug_features.size(0), 1)
+        text_direction /= text_direction.norm(dim=-1, keepdim=True)
+
+        loss_calc = (1 - torch.cosine_similarity(img_direction, text_direction, dim=1))
+        loss_calc[loss_calc < PATCH_THRESHOLD] = 0.0        
+        loss_patch += loss_calc.mean()
+
+        glob_features = clip_model.get_image_features(pixel_values=clip_normalize(target))
+        glob_features /= (glob_features.clone().norm(dim=-1, keepdim=True))
+        glob_direction = glob_features - img_features
+        glob_direction /= glob_direction.clone().norm(dim=-1, keepdim=True)
+
+        loss_glob = (1 - torch.cosine_similarity(glob_direction, text_direction, dim=1)).mean()
+
+        loss_tv = L_TV * get_image_prior_losses(target)
+
+        total_loss = (L_PATCH * loss_patch) + (L_DIR * loss_glob) + (L_CONTENT * content_loss) + loss_tv
+        total_loss_epoch.append(total_loss)
+
+        optimizer.zero_grad()
+        total_loss.backward()
+        optimizer.step()
+    
+    # Post-process output
+    output_image = target.clone().detach()
+    output_image = torch.clamp(output_image, 0, 1)
+    output_image = adjust_contrast(output_image, 1.5)
+    
+    # Convert tensor back to PIL Image
+    output_image = output_image.squeeze(0).cpu()
+    output_image = transforms.ToPILImage()(output_image)
+    
+    return output_image

util.py

@@ -129,8 +129,7 @@ def clip_normalize(data):
     norm_data = (resized - mean) / std
     return norm_data
 
-def load_image(img_path):
-    image = Image.open(img_path)
+def load_image(image):
     image = image.resize((IMG_SIZE, IMG_SIZE))
     transform = transforms.Compose([transforms.ToTensor()])
     return transform(image)[:3, :, :].unsqueeze(0)