new-cnn (#1)

- trying new cnn (3601849835ba31be17f768a180841e562399333c)

app.py

@@ -4,21 +4,39 @@ import clip
 from torchvision import transforms
 from style_net import StyleNet
 from dataset import FlickrStreamer
+from util import DEVICE
+from clip_styler import ClipStyler
 
 # --- 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. Editing Model
-style_net = StyleNet().to(DEVICE)
+# A. Retrieval Model (CLIP)
+clip_model, _ = clip.load("ViT-B/32", device=DEVICE)
 
-# B. Retrieval Model (CLIP)
-clip_model, preprocess = clip.load("ViT-B/32", device=DEVICE)
+# B. Editing Model (CLIP-guided optimizer)
+clip_styler = ClipStyler(device=DEVICE, num_epochs=UI_EPOCHS)
 
 # 2. LOAD SEARCH INDEX
 print("Loading Search Index...")
@@ -60,32 +78,25 @@ def search_flickr(text_query):
             
         return results
 
-# --- FUNCTION: EDIT ---
-tf_edit = transforms.Compose([
-    transforms.Resize((256, 256)),
-    transforms.ToTensor()
-])
-
 def edit_image(image, style_choice):
-    if image is None: return None
-    
-    weights_map = {
-        "Sketch": "model_sketch.pth",
-        "Van Gogh": "model_painting.pth",
-        "Cyberpunk": "model_cyberpunk.pth"
-    }
-    
+    if image is None or style_choice is None:
+        return None
+
+    preset = STYLE_PRESETS.get(style_choice)
+    if preset is None:
+        return image
+
     try:
-        style_net.load_state_dict(torch.load(weights_map.get(style_choice), map_location=DEVICE))
-    except Exception:
+        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}")
         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
 
 # --- THE UI ---
 with gr.Blocks() as demo:

clip_styler.py

@@ -0,0 +1,212 @@
+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 transformers import CLIPModel, AutoTokenizer
+
+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:
+    """
+    Wraps the CLIP-guided optimization loop from main.py so it can be reused
+    by both the CLI entrypoint and the Gradio application.
+    """
+
+    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)
+        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
+

cnn1.py

@@ -0,0 +1,251 @@
+import torch
+import torch.nn as nn
+
+
+class FiLMLayer(nn.Module):
+    """
+    Feature-wise linear modulation module that conditions convolutional activations
+    on an external style embedding (e.g., a CLIP text embedding).
+    """
+
+    def __init__(self, num_channels: int, cond_dim: int, hidden_dim: int = 256):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(cond_dim),
+            nn.Linear(cond_dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, num_channels * 2),
+        )
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        gamma, beta = self.net(cond).chunk(2, dim=1)
+        gamma = gamma.unsqueeze(-1).unsqueeze(-1)
+        beta = beta.unsqueeze(-1).unsqueeze(-1)
+        return x * (1 + gamma) + beta
+
+
+class DoubleConv(nn.Module):
+    """Two consecutive conv-batchnorm-gelu blocks with optional FiLM conditioning."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        cond_dim: int | None = None,
+        film_hidden_dim: int = 256,
+    ):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.GELU(),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.GELU(),
+        )
+        self.film = (
+            FiLMLayer(out_channels, cond_dim, hidden_dim=film_hidden_dim)
+            if cond_dim is not None
+            else None
+        )
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor | None = None) -> torch.Tensor:
+        x = self.conv(x)
+        if self.film is not None:
+            if cond is None:
+                raise ValueError("Style embedding is required for FiLM conditioning.")
+            x = self.film(x, cond)
+        return x
+
+
+class DownBlock(nn.Module):
+    """Down-sampling block used in the encoder path."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        cond_dim: int | None = None,
+        film_hidden_dim: int = 256,
+    ):
+        super().__init__()
+        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.conv = DoubleConv(
+            in_channels, out_channels, cond_dim, film_hidden_dim=film_hidden_dim
+        )
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor | None = None) -> torch.Tensor:
+        x = self.pool(x)
+        return self.conv(x, cond)
+
+
+class UpBlock(nn.Module):
+    """Up-sampling block with skip connections from the encoder path."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        skip_channels: int,
+        cond_dim: int | None = None,
+        bilinear: bool = True,
+        film_hidden_dim: int = 256,
+    ):
+        super().__init__()
+        if bilinear:
+            self.up = nn.Sequential(
+                nn.Upsample(scale_factor=2, mode="bilinear", align_corners=True),
+                nn.Conv2d(in_channels, in_channels // 2, kernel_size=1),
+            )
+        else:
+            self.up = nn.ConvTranspose2d(
+                in_channels, in_channels // 2, kernel_size=2, stride=2
+            )
+        self.conv = DoubleConv(
+            in_channels // 2 + skip_channels,
+            skip_channels,
+            cond_dim,
+            film_hidden_dim=film_hidden_dim,
+        )
+
+    def forward(
+        self, x: torch.Tensor, skip: torch.Tensor, cond: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        x = self.up(x)
+        diff_y = skip.size(2) - x.size(2)
+        diff_x = skip.size(3) - x.size(3)
+        if diff_y != 0 or diff_x != 0:
+            x = nn.functional.pad(
+                x,
+                [
+                    diff_x // 2,
+                    diff_x - diff_x // 2,
+                    diff_y // 2,
+                    diff_y - diff_y // 2,
+                ],
+            )
+        x = torch.cat([skip, x], dim=1)
+        return self.conv(x, cond)
+
+
+class OutConv(nn.Module):
+    """Final projection into the RGB space."""
+
+    def __init__(self, in_channels: int, out_channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+class UNet(nn.Module):
+    """
+    Lightweight encoder-decoder network for CLIP-guided, text-prompted style transfer.
+
+    The network takes a content image and optionally a CLIP text embedding that
+    modulates intermediate activations through FiLM layers so that the decoded
+    image aligns with the target style semantics in CLIP space.
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        base_channels: int = 16,
+        num_layers: int = 4,
+        text_dim: int = None,
+        bilinear: bool = True,
+        film_hidden_dim: int = 256,
+    ):
+        super().__init__()
+
+        if num_layers < 2:
+            raise ValueError("num_layers must be >= 2")
+
+        self.cond_dim = text_dim
+        self.style_mapper = (
+            nn.Sequential(
+                nn.LayerNorm(text_dim),
+                nn.Linear(text_dim, text_dim),
+                nn.GELU(),
+                nn.Linear(text_dim, text_dim),
+            )
+            if text_dim is not None
+            else None
+        )
+
+        channels = [base_channels * (2**i) for i in range(num_layers)]
+
+        self.inc = DoubleConv(
+            in_channels,
+            channels[0],
+            self.cond_dim,
+            film_hidden_dim=film_hidden_dim,
+        )
+        self.downs = nn.ModuleList()
+        for idx in range(num_layers - 1):
+            self.downs.append(
+                DownBlock(
+                    channels[idx],
+                    channels[idx + 1],
+                    self.cond_dim,
+                    film_hidden_dim=film_hidden_dim,
+                )
+            )
+
+        self.bottleneck = DoubleConv(
+            channels[-1],
+            channels[-1] * 2,
+            self.cond_dim,
+            film_hidden_dim=film_hidden_dim,
+        )
+
+        self.ups = nn.ModuleList()
+        prev_channels = channels[-1] * 2
+        for skip_ch in reversed(channels):
+            self.ups.append(
+                UpBlock(
+                    prev_channels,
+                    skip_ch,
+                    self.cond_dim,
+                    bilinear=bilinear,
+                    film_hidden_dim=film_hidden_dim,
+                )
+            )
+            prev_channels = skip_ch
+
+        self.outc = OutConv(channels[0], out_channels)
+        self.activation = nn.Tanh()
+
+    def _prepare_condition(self, text_embedding: torch.Tensor | None) -> torch.Tensor | None:
+        if self.cond_dim is None:
+            return None
+        if text_embedding is None:
+            raise ValueError(
+                "text_embedding must be provided when the model is configured for conditioning."
+            )
+        if text_embedding.dim() != 2 or text_embedding.size(1) != self.cond_dim:
+            raise ValueError(
+                f"text_embedding must have shape [batch, {self.cond_dim}] but got {text_embedding.shape}."
+            )
+        return self.style_mapper(text_embedding) if self.style_mapper else text_embedding
+
+    def forward(
+        self, x: torch.Tensor, text_embedding: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        cond = self._prepare_condition(text_embedding)
+
+        skip_connections = []
+        x = self.inc(x, cond)
+        skip_connections.append(x)
+
+        for down in self.downs:
+            x = down(x, cond)
+            skip_connections.append(x)
+
+        x = self.bottleneck(x, cond)
+
+        for up, skip in zip(self.ups, reversed(skip_connections)):
+            x = up(x, skip, cond)
+
+        x = self.outc(x)
+        return self.activation(x)

requirements.txt

@@ -5,4 +5,5 @@ regex
 tqdm
 deeplake<4
 gradio
-git+https://github.com/openai/CLIP.git
+git+https://github.com/openai/CLIP.git
+transformers

style_net.py

@@ -1,63 +1,63 @@
-import torch
-import torch.nn as nn
-
-class ResidualBlock(nn.Module):
-    def __init__(self, channels):
-        super(ResidualBlock, self).__init__()
-        self.conv = nn.Sequential(
-            nn.Conv2d(channels, channels, 3, padding=1),
-            nn.InstanceNorm2d(channels),
-            nn.ReLU(),
-            nn.Conv2d(channels, channels, 3, padding=1),
-            nn.InstanceNorm2d(channels)
-        )
-
-    def forward(self, x):
-        return x + self.conv(x)
-    
-class UpsampleConvLayer(nn.Module):
-    def __init__(self, in_channels, out_channels, kernel_size, stride):
-        super().__init__()
-        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
-        # Reflection padding reduces "border" artifacts
-        self.reflection_pad = nn.ReflectionPad2d(kernel_size // 2)
-        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=1)
-
-    def forward(self, x):
-        x = self.upsample(x)
-        x = self.reflection_pad(x)
-        return self.conv(x)
-
-class StyleNet(nn.Module):
-    def __init__(self):
-        super(StyleNet, self).__init__()
-        # Encoder
-        self.encoder = nn.Sequential(
-            nn.ReflectionPad2d(4),           # Pad first...
-            nn.Conv2d(3, 32, 9, padding=0),  # ...then Conv (padding=0)
-            nn.InstanceNorm2d(32), nn.ReLU(),
-            
-            nn.ReflectionPad2d(1),
-            nn.Conv2d(32, 64, 3, stride=2, padding=0),
-            nn.InstanceNorm2d(64), nn.ReLU(),
-            
-            nn.ReflectionPad2d(1),
-            nn.Conv2d(64, 128, 3, stride=2, padding=0),
-            nn.InstanceNorm2d(128), nn.ReLU()
-        )
-        # Bottleneck
-        self.bottleneck = nn.Sequential(*[ResidualBlock(128) for _ in range(3)])
-        # Decoder
-        self.decoder = nn.Sequential(
-            UpsampleConvLayer(128, 64, kernel_size=3, stride=1),
-            nn.InstanceNorm2d(64), nn.ReLU(),
-            UpsampleConvLayer(64, 32, kernel_size=3, stride=1),
-            nn.InstanceNorm2d(32), nn.ReLU(),
-            nn.Conv2d(32, 3, 9, padding=4),
-            nn.Sigmoid()
-        )
-
-    def forward(self, x):
-        features = self.encoder(x)
-        features = self.bottleneck(features)
+import torch
+import torch.nn as nn
+
+class ResidualBlock(nn.Module):
+    def __init__(self, channels):
+        super(ResidualBlock, self).__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(channels, channels, 3, padding=1),
+            nn.InstanceNorm2d(channels),
+            nn.ReLU(),
+            nn.Conv2d(channels, channels, 3, padding=1),
+            nn.InstanceNorm2d(channels)
+        )
+
+    def forward(self, x):
+        return x + self.conv(x)
+    
+class UpsampleConvLayer(nn.Module):
+    def __init__(self, in_channels, out_channels, kernel_size, stride):
+        super().__init__()
+        self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)
+        # Reflection padding reduces "border" artifacts
+        self.reflection_pad = nn.ReflectionPad2d(kernel_size // 2)
+        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride=1)
+
+    def forward(self, x):
+        x = self.upsample(x)
+        x = self.reflection_pad(x)
+        return self.conv(x)
+
+class StyleNet(nn.Module):
+    def __init__(self):
+        super(StyleNet, self).__init__()
+        # Encoder
+        self.encoder = nn.Sequential(
+            nn.ReflectionPad2d(4),           # Pad first...
+            nn.Conv2d(3, 32, 9, padding=0),  # ...then Conv (padding=0)
+            nn.InstanceNorm2d(32), nn.ReLU(),
+            
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(32, 64, 3, stride=2, padding=0),
+            nn.InstanceNorm2d(64), nn.ReLU(),
+            
+            nn.ReflectionPad2d(1),
+            nn.Conv2d(64, 128, 3, stride=2, padding=0),
+            nn.InstanceNorm2d(128), nn.ReLU()
+        )
+        # Bottleneck
+        self.bottleneck = nn.Sequential(*[ResidualBlock(128) for _ in range(3)])
+        # Decoder
+        self.decoder = nn.Sequential(
+            UpsampleConvLayer(128, 64, kernel_size=3, stride=1),
+            nn.InstanceNorm2d(64), nn.ReLU(),
+            UpsampleConvLayer(64, 32, kernel_size=3, stride=1),
+            nn.InstanceNorm2d(32), nn.ReLU(),
+            nn.Conv2d(32, 3, 9, padding=4),
+            nn.Sigmoid()
+        )
+
+    def forward(self, x):
+        features = self.encoder(x)
+        features = self.bottleneck(features)
         return self.decoder(features)

util.py

@@ -0,0 +1,168 @@
+from PIL import Image
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torchvision import transforms, models
+
+IMG_MEAN = [0.485, 0.456, 0.406]
+IMG_STD = [0.229, 0.224, 0.225]
+CLIP_MEAN = [0.48145466, 0.4578275, 0.40821073]
+CLIP_STD = [0.26862954, 0.26130258, 0.27577711]
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+LR = 5e-4
+CROP_SIZE = 128
+RESIZE = 224
+NUM_EPOCHS = 200
+NUM_CROPS = 64
+PATCH_THRESHOLD = 0.7
+L_TV = 2e-3
+L_PATCH = 9000
+L_DIR = 500
+L_CONTENT = 150
+IMG_SIZE = 512
+# copied from openai clip
+IMAGENET_TEMPLATES = [
+    "a bad photo of a {}.",
+    "a photo of many {}.",
+    "a sculpture of a {}.",
+    "a photo of the hard to see {}.",
+    "a low resolution photo of the {}.",
+    "a rendering of a {}.",
+    "graffiti of a {}.",
+    "a bad photo of the {}.",
+    "a cropped photo of the {}.",
+    "a tattoo of a {}.",
+    "the embroidered {}.",
+    "a photo of a hard to see {}.",
+    "a bright photo of a {}.",
+    "a photo of a clean {}.",
+    "a photo of a dirty {}.",
+    "a dark photo of the {}.",
+    "a drawing of a {}.",
+    "a photo of my {}.",
+    "the plastic {}.",
+    "a photo of the cool {}.",
+    "a close-up photo of a {}.",
+    "a black and white photo of the {}.",
+    "a painting of the {}.",
+    "a painting of a {}.",
+    "a pixelated photo of the {}.",
+    "a sculpture of the {}.",
+    "a bright photo of the {}.",
+    "a cropped photo of a {}.",
+    "a plastic {}.",
+    "a photo of the dirty {}.",
+    "a jpeg corrupted photo of a {}.",
+    "a blurry photo of the {}.",
+    "a photo of the {}.",
+    "a good photo of the {}.",
+    "a rendering of the {}.",
+    "a {} in a video game.",
+    "a photo of one {}.",
+    "a doodle of a {}.",
+    "a close-up photo of the {}.",
+    "a photo of a {}.",
+    "the origami {}.",
+    "the {} in a video game.",
+    "a sketch of a {}.",
+    "a doodle of the {}.",
+    "a origami {}.",
+    "a low resolution photo of a {}.",
+    "the toy {}.",
+    "a rendition of the {}.",
+    "a photo of the clean {}.",
+    "a photo of a large {}.",
+    "a rendition of a {}.",
+    "a photo of a nice {}.",
+    "a photo of a weird {}.",
+    "a blurry photo of a {}.",
+    "a cartoon {}.",
+    "art of a {}.",
+    "a sketch of the {}.",
+    "a embroidered {}.",
+    "a pixelated photo of a {}.",
+    "itap of the {}.",
+    "a jpeg corrupted photo of the {}.",
+    "a good photo of a {}.",
+    "a plushie {}.",
+    "a photo of the nice {}.",
+    "a photo of the small {}.",
+    "a photo of the weird {}.",
+    "the cartoon {}.",
+    "art of the {}.",
+    "a drawing of the {}.",
+    "a photo of the large {}.",
+    "a black and white photo of a {}.",
+    "the plushie {}.",
+    "a dark photo of a {}.",
+    "itap of a {}.",
+    "graffiti of the {}.",
+    "a toy {}.",
+    "itap of my {}.",
+    "a photo of a cool {}.",
+    "a photo of a small {}.",
+    "a tattoo of the {}.",
+]
+
+def get_mean(mean_dist):
+    mean = torch.tensor(mean_dist).to(DEVICE)
+    return mean.view(1, -1, 1, 1)
+
+def get_std(std_dist):
+    std = torch.tensor(std_dist).to(DEVICE)
+    return std.view(1, -1, 1, 1)
+
+def normalize(data):
+    mean = get_mean(IMG_MEAN)
+    std = get_std(IMG_STD)
+
+    norm_data = (data - mean) / std
+    return norm_data
+
+def clip_normalize(data):
+    resized = nn.functional.interpolate(data, size=RESIZE, mode='bicubic')
+    mean = get_mean(CLIP_MEAN)
+    std = get_std(CLIP_STD)
+    norm_data = (resized - mean) / std
+    return norm_data
+
+def load_image(img_path):
+    image = Image.open(img_path)
+    image = image.resize((IMG_SIZE, IMG_SIZE))
+    transform = transforms.Compose([transforms.ToTensor()])
+    return transform(image)[:3, :, :].unsqueeze(0)
+
+def get_features(image, vgg19):
+    # uses vgg19 model to extract content features
+    layers = {'0': 'conv1_1',
+              '5': 'conv2_1',
+              '10': 'conv3_1',
+              '19': 'conv4_1',
+              '21': 'conv4_2',
+              '28': 'conv5_1',
+              '31': 'conv5_2'
+              }
+    features = {}
+    x = image
+    for name, layer in vgg19._modules.items():
+        x = layer(x)
+        if name in layers:
+            features[layers[name]] = x
+
+    return features
+
+def prompt_ensemble(prompt):
+    return [template.format(prompt) for template in IMAGENET_TEMPLATES]
+
+def get_image_prior_losses(target):
+    diff1 = target[:, :, :, :-1] - target[:, :, :, 1:]
+    diff2 = target[:, :, :-1, :] - target[:, :, 1:, :]
+    diff3 = target[:, :, 1:, :-1] - target[:, :, :-1, 1:]
+    diff4 = target[:, :, :-1, :-1] - target[:, :, 1:, 1:]
+
+    loss_var_l2 = torch.norm(diff1) + torch.norm(diff2) + torch.norm(diff3) + torch.norm(diff4)
+    
+    return loss_var_l2