clip_styler.py

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