Spaces:
Sleeping
Sleeping
File size: 2,599 Bytes
5a14c00 3048202 5a14c00 3048202 52323ef 5a14c00 52323ef 3048202 52323ef 3048202 52323ef 3048202 52323ef 5a14c00 3048202 5a14c00 3048202 52323ef 3048202 52323ef 3048202 5a14c00 52323ef 3048202 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | 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
import torch.nn.functional as F
def load_models(weights_path):
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(DEVICE)
tokenizer = AutoTokenizer.from_pretrained("openai/clip-vit-base-patch32")
vgg19 = models.vgg19(pretrained=True).features.to(DEVICE)
for param in vgg19.parameters():
param.requires_grad_(False)
cnn_model = UNet(text_dim=512).to(DEVICE)
cnn_model.load_state_dict(torch.load(weights_path, map_location=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, weights_path, 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(weights_path)
cnn_model.eval()
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)
with torch.inference_mode():
target = cnn_model(img, text_embedding=text_features).to(DEVICE)
# 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 |