using inference

app.py

@@ -67,19 +67,19 @@ tf_edit = transforms.Compose([
     transforms.ToTensor()
 ])
 
-def edit_image(image, style_choice, num_steps):
+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"
-    }
     detail_prompt = {
         "Sketch": "charcoal sketch",
         "Van Gogh": "Van Gogh painting",
         "Cyberpunk": "neon cyberpunk",
     }
+    WEIGHTS_MAP = {
+        "Sketch": "unet_charcoal-sketch.pth",
+        "Van Gogh": "unet_van-gogh-painting.pth",
+        "Cyberpunk": "unet_neon-cyberpunk.pth"
+    }
     '''
     try:
         style_net.load_state_dict(torch.load(weights_map.get(style_choice), map_location=DEVICE))
@@ -93,7 +93,7 @@ def edit_image(image, style_choice, num_steps):
     output_img = transforms.ToPILImage()(output_tensor.squeeze().cpu())
     return output_img
     '''
-    return style_transfer(image, detail_prompt[style_choice], num_steps=num_steps)
+    return style_transfer(image, detail_prompt[style_choice], weights_path=WEIGHTS_MAP[style_choice])
 
 # --- THE UI ---
 with gr.Blocks() as demo:
@@ -115,18 +115,11 @@ 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, steps_slider], outputs=im_out)
+        btn_edit.click(edit_image, inputs=[im_in, style], outputs=im_out)
 
 if __name__ == "__main__":
     demo.launch(share=True)

clip_styler.py

@@ -5,29 +5,18 @@ 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
 
-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")
+def load_models(weights_path):
     
-    if vgg19 is None:
-        vgg19 = models.vgg19(pretrained=True).features.to(DEVICE)
-        for param in vgg19.parameters():
-            param.requires_grad_(False)
+    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)
     
-    if cnn_model is None:
-        cnn_model = UNet(text_dim=512).to(DEVICE)
+    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
 
@@ -42,7 +31,7 @@ def reset_cnn_model():
                 if layer.bias is not None:
                     torch.nn.init.zeros_(layer.bias)
 
-def style_transfer(img, prompt, num_steps=NUM_EPOCHS, source="a Photo"):
+def style_transfer(img, prompt, weights_path, source="a Photo"):
     """
     Apply style transfer to an uploaded image
     
@@ -60,33 +49,9 @@ def style_transfer(img, prompt, num_steps=NUM_EPOCHS, source="a Photo"):
     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)
+    clip_model, tokenizer, vgg19, cnn_model = load_models(weights_path)
 
-    # 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)
+    cnn_model.eval()
 
     with torch.no_grad():
         edited_text = prompt_ensemble(prompt)
@@ -94,65 +59,9 @@ def style_transfer(img, prompt, num_steps=NUM_EPOCHS, source="a Photo"):
         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)
+        
+    with torch.inference_mode():
         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()