app.py

import zipfile
import subprocess
def unzip_content():
    try:
        print("Attempting to unzip content using Python...")
        with zipfile.ZipFile('./content.zip', 'r') as zip_ref:
            zip_ref.extractall('.')
    except Exception as e:
        print(f"Python unzip failed: {str(e)}")
        try:
            print("Attempting to unzip content using system command...")
            subprocess.run(['unzip', '-o', './content.zip'], check=True)
        except Exception as e:
            print(f"System unzip failed: {str(e)}")
            raise Exception("Failed to unzip content using both methods")
    print("Content successfully unzipped!")
import subprocess
import gradio as gr
import numpy as np
import torch
import torchvision
import torchvision.transforms
import torchvision.transforms.functional
import PIL
import matplotlib.pyplot as plt
import yaml
from omegaconf import OmegaConf
from CLIP import clip
import os
import sys
import tempfile
import io
from pathlib import Path
from PIL import Image
import cv2
import imageio

# Add taming transformers path
taming_path = os.path.join(os.getcwd(), 'taming-transformers')
sys.path.append(taming_path)
from taming.models.vqgan import VQModel

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")



def create_video(images_list, video_name='morphing_video.mp4'):
    """Create video from a list of image tensors"""
    if not images_list:
        print("No images provided.")
        return None

    with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_video:
        video_writer = imageio.get_writer(temp_video.name, fps=10)
        
        for img_tensor in images_list:
            # Convert tensor to numpy array
            img = img_tensor.cpu().numpy().transpose((1, 2, 0))
            img = (img * 255).astype('uint8')
            video_writer.append_data(img)

        video_writer.close()
        return temp_video.name

def save_from_tensors(tensor):
    """Process tensor and return the processed version"""
    img = tensor.clone()
    img = img.mul(255).byte()
    img = img.cpu().numpy().transpose((1, 2, 0))
    return img

def norm_data(data):
    return (data.clip(-1, 1) + 1) / 2

def setup_clip_model():
    model, _ = clip.load('ViT-B/32', jit=False)
    model.eval().to(device)
    return model

def setup_vqgan_model(config_path, checkpoint_path):
    config = OmegaConf.load(config_path)
    model = VQModel(**config.model.params)
    state_dict = torch.load(checkpoint_path, map_location="cpu")["state_dict"]
    model.load_state_dict(state_dict, strict=False)
    return model.eval().to(device)

def generator(x, model):
    x = model.post_quant_conv(x)
    x = model.decoder(x)
    return x

def encode_text(text, clip_model):
    t = clip.tokenize(text).to(device)
    return clip_model.encode_text(t).detach().clone()

def create_encoding(include, exclude, extras, clip_model):
    include_enc = [encode_text(text, clip_model) for text in include]
    exclude_enc = [encode_text(text, clip_model) for text in exclude]
    extras_enc = [encode_text(text, clip_model) for text in extras]
    return include_enc, exclude_enc, extras_enc

def create_crops(img, num_crops=32, size1=225, noise_factor=0.05):
    aug_transform = torch.nn.Sequential(
        torchvision.transforms.RandomHorizontalFlip(),
        torchvision.transforms.RandomAffine(30, translate=(0.1, 0.1), fill=0)
    ).to(device)

    p = size1 // 2
    img = torch.nn.functional.pad(img, (p, p, p, p), mode='constant', value=0)
    img = aug_transform(img)

    crop_set = []
    for _ in range(num_crops):
        gap1 = int(torch.normal(1.2, .3, ()).clip(.43, 1.9) * size1)
        offsetx = torch.randint(0, int(size1 * 2 - gap1), ())
        offsety = torch.randint(0, int(size1 * 2 - gap1), ())
        crop = img[:, :, offsetx:offsetx + gap1, offsety:offsety + gap1]
        crop = torch.nn.functional.interpolate(crop, (224, 224), mode='bilinear', align_corners=True)
        crop_set.append(crop)

    img_crops = torch.cat(crop_set, 0)
    randnormal = torch.randn_like(img_crops, requires_grad=False)
    randstotal = torch.rand((img_crops.shape[0], 1, 1, 1)).to(device)
    img_crops = img_crops + noise_factor * randstotal * randnormal

    return img_crops

def optimize_result(params, prompt, vqgan_model, clip_model, w1, w2, extras_enc, exclude_enc):
    alpha = 1
    beta = 0.5
    out = generator(params, vqgan_model)
    out = norm_data(out)
    out = create_crops(out)
    out = torchvision.transforms.Normalize((0.48145466, 0.4578275, 0.40821073),
                                         (0.26862954, 0.26130258, 0.27577711))(out)

    img_enc = clip_model.encode_image(out)
    final_enc = w1 * prompt + w2 * extras_enc[0]
    final_text_include_enc = final_enc / final_enc.norm(dim=-1, keepdim=True)
    final_text_exclude_enc = exclude_enc[0]

    main_loss = torch.cosine_similarity(final_text_include_enc, img_enc, dim=-1)
    penalize_loss = torch.cosine_similarity(final_text_exclude_enc, img_enc, dim=-1)

    return -alpha * main_loss.mean() + beta * penalize_loss.mean()

def optimize(params, optimizer, prompt, vqgan_model, clip_model, w1, w2, extras_enc, exclude_enc):
    loss = optimize_result(params, prompt, vqgan_model, clip_model, w1, w2, extras_enc, exclude_enc)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss

def training_loop(params, optimizer, include_enc, exclude_enc, extras_enc, vqgan_model, clip_model, w1, w2,
                 total_iter=200, show_step=1):
    res_img = []
    res_z = []

    for prompt in include_enc:
        for it in range(total_iter):
            loss = optimize(params, optimizer, prompt, vqgan_model, clip_model, w1, w2, extras_enc, exclude_enc)

            if it >= 0 and it % show_step == 0:
                with torch.no_grad():
                    generated = generator(params, vqgan_model)
                    new_img = norm_data(generated[0].to(device))
                res_img.append(new_img)
                res_z.append(params.clone().detach())
                print(f"loss: {loss.item():.4f}\nno. of iteration: {it}")

        torch.cuda.empty_cache()
    return res_img, res_z

def generate_art(include_text, exclude_text, extras_text, num_iterations):
    try:
        # Process the input prompts
        include = [x.strip() for x in include_text.split(',')]
        exclude = [x.strip() for x in exclude_text.split(',')]
        extras = [x.strip() for x in extras_text.split(',')]
        
        w1, w2 = 1.0, 0.9

        # Setup models
        clip_model = setup_clip_model()
        vqgan_model = setup_vqgan_model("./models/vqgan_imagenet_f16_16384/configs/model.yaml",
                                       "./models/vqgan_imagenet_f16_16384/checkpoints/last.ckpt")

        # Parameters
        learning_rate = 0.1
        batch_size = 1
        wd = 0.1
        size1, size2 = 225, 400

        # Initialize parameters
        initial_image = PIL.Image.open('./gradient1.png')
        initial_image = initial_image.resize((size2, size1))
        initial_image = torchvision.transforms.ToTensor()(initial_image).unsqueeze(0).to(device)

        with torch.no_grad():
            z, _, _ = vqgan_model.encode(initial_image)

        params = torch.nn.Parameter(z).to(device)
        optimizer = torch.optim.AdamW([params], lr=learning_rate, weight_decay=wd)
        params.data = params.data * 0.6 + torch.randn_like(params.data) * 0.4

        # Encode prompts
        include_enc, exclude_enc, extras_enc = create_encoding(include, exclude, extras, clip_model)

        # Run training loop
        res_img, res_z = training_loop(params, optimizer, include_enc, exclude_enc, extras_enc,
                                     vqgan_model, clip_model, w1, w2, total_iter=num_iterations)

        # Create video directly from tensors
        video_path = create_video(res_img)
        
        return video_path
        
    except Exception as e:
        raise e

def gradio_interface(include_text, exclude_text, extras_text, num_iterations):
    try:
        video_path = generate_art(include_text, exclude_text, extras_text, int(num_iterations))
        return video_path
    except Exception as e:
        return f"An error occurred: {str(e)}"

# Try to unzip content at startup
try:
    unzip_content()
except Exception as e:
    print(f"Warning: Could not unzip content: {str(e)}")

# Define and launch the Gradio app
iface = gr.Interface(
    fn=gradio_interface,
    inputs=[
        gr.Textbox(label="Include Prompts (comma-separated)", 
                  value="desert, heavy rain, cactus"),
        gr.Textbox(label="Exclude Prompts (comma-separated)", 
                  value="confusing, blurry"),
        gr.Textbox(label="Extra Style Prompts (comma-separated)", 
                  value="desert, clear, detailed, beautiful, good shape, detailed"),
        gr.Number(label="Number of Iterations", 
                 value=200, minimum=1, maximum=1000)
    ],
    outputs=gr.Video(label="Generated Morphing Video", format="mp4", autoplay=True),
    title="VQGAN-CLIP Art Generator",
    css="allow",
    allow_flagging="never",
    description = """
Generate artistic videos using VQGAN-CLIP. Enter your prompts separated by commas and adjust the number of iterations. The model will generate a morphing video based on your inputs.
Note: This application requires GPU access. Please either:
1. Use the Colab notebook available at https://github.com/SanshruthR/VQGAN-CLIP
2. Clone this space and enable GPU in your personal copy.
""")

if __name__ == "__main__":
    print("Checking GPU availability:", "GPU AVAILABLE" if torch.cuda.is_available() else "NO GPU FOUND")
    iface.launch()