Spaces:
Sleeping
Sleeping
| import torch | |
| 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 | |
| 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"): | |
| """ | |
| 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 | |
| """ | |
| 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) | |
| 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 |