app.py

import os
import io
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import transforms
from PIL import Image, ImageTk, ImageFilter
import numpy as np
import gradio as gr
from huggingface_hub import hf_hub_download


# --- Hyperparameters ---
image_size = 64
latent_dim = 128
model_repo_id = "elapt1c/catGen"
model_filename = "model.pth"
#model_path = 'model.pth'  # Relative path within the space. Assumed it will be in the root
generated_images_folder = 'generated_images'


# --- VAE Model --- (Simplified VAE - MATCHING TRAINING CODE)
class VAE(nn.Module):
    def __init__(self, latent_dim):
        super(VAE, self).__init__()

        # Encoder - MATCHING TRAINING CODE ARCHITECTURE
        self.encoder_conv = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=4, stride=2, padding=1),  # Increased initial channels
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(32, 64, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1),  # Increased final channels
            nn.LeakyReLU(0.2, inplace=True),
        )
        self.encoder_fc_mu = nn.Linear(512 * 2 * 2, latent_dim)
        self.encoder_fc_logvar = nn.Linear(512 * 2 * 2, latent_dim)

        # Decoder - MATCHING TRAINING CODE ARCHITECTURE
        self.decoder_fc = nn.Linear(latent_dim, 512 * 2 * 2)
        self.decoder_conv = nn.Sequential(
            nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),
            nn.LeakyReLU(0.2, inplace=True),
            nn.ConvTranspose2d(32, 3, kernel_size=4, stride=2, padding=1),
            nn.Sigmoid()
        )

    def encode(self, x):
        h = self.encoder_conv(x)
        h = h.view(h.size(0), -1)
        mu = self.encoder_fc_mu(h)
        logvar = self.encoder_fc_logvar(h)
        return mu, logvar

    def decode(self, z):
        z = self.decoder_fc(z)
        z = z.view(z.size(0), 512, 2, 2)  # Corrected view shape to 512 channels
        reconstructed_image = self.decoder_conv(z)
        return reconstructed_image

    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std

    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        reconstructed_image = self.decode(z)
        return reconstructed_image, mu, logvar


# --- Helper Functions ---
def load_model(device, repo_id, filename):
    try:
        model_path = hf_hub_download(repo_id=repo_id, filename=filename)
    except Exception as e:
        print(f"Error downloading model from Hugging Face Hub: {e}")
        return None

    vae_model = VAE(latent_dim=latent_dim).to(device)  # Plain VAE model

    try:
        checkpoint = torch.load(model_path, map_location=device)  # Load checkpoint dict
    except FileNotFoundError:
        print(f"Error: Model file not found at {model_path}. This should not happen after downloading.")
        return None

    new_state_dict = {}  # Create a new dictionary for modified keys
    for key, value in checkpoint.items():
        new_key = key.replace('_orig_mod.', '')  # Remove "_orig_mod." prefix
        new_state_dict[new_key] = value  # Add to new dict with modified key

    vae_model.load_state_dict(new_state_dict)  # Load state_dict with modified keys
    print(f"====> Loaded existing model from {model_path} (handling Torch Compile state_dict)")
    return vae_model


def preprocess_image(image):
    try:
        transform = transforms.Compose([
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
        ])
        image = transform(image).unsqueeze(0)
        return image
    except Exception as e:
        print(f"Failed to preprocess image: {e}")
        return None


def generate_single_image(model, device):
    try:
        model.eval()
        with torch.no_grad():
            sample_z = torch.randn(1, latent_dim).to(device)
            generated_image = model.decode(sample_z)  # Use simple VAE decode
            img = generated_image.cpu().detach().numpy()
            output = (img[0] * 255).transpose(1, 2, 0).astype(np.uint8)
            image = Image.fromarray(output)  # save from random image
            return image  # use the image
    except Exception as e:
        print(f"Image generation failed: {e}")
        return None


def generate_from_base_image(model, device, base_image, noise_scale=0.1):
    try:
        model.eval()
        with torch.no_grad():
            processed_image = preprocess_image(base_image)  # Process base image
            if processed_image is None:
                return None

            processed_image = processed_image.to(device)  # to device
            mu, logvar = model.encode(processed_image)  # encode
            latent_vector = model.reparameterize(mu, logvar)  # reparameterize

            noise = torch.randn_like(latent_vector) * noise_scale  # add noise
            latent_vector = latent_vector + noise  # combine

            generated_image = model.decode(latent_vector)  # Use simple VAE decode
            img = generated_image.cpu().detach().numpy()
            output = (img[0] * 255).transpose(1, 2, 0).astype(np.uint8)
            output_image = Image.fromarray(output)  # save from
            return output_image

    except Exception as e:
        print(f"Seed image generation failed: {e}")
        return None



# --- Gradio Interface ---
def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    vae_model = load_model(device, model_repo_id, model_filename)
    if vae_model is None:
        return  # Exit if model loading fails

    def generate_single():
        img = generate_single_image(vae_model, device)
        if img:
            return img
        else:
            return "Image generation failed. Check console for errors."

    def generate_from_seed(seed_image):
        if seed_image is None:
            return "Please upload a seed image."

        img = generate_from_base_image(vae_model, device, seed_image)
        if img:
            return img
        else:
            return "Image generation from seed failed. Check console for errors."


    with gr.Blocks() as demo:
        gr.Markdown("# VAE Image Generator")

        with gr.Tab("Generate Single Image"):
            single_button = gr.Button("Generate Random Image")
            single_output = gr.Image()
            single_button.click(generate_single, inputs=[], outputs=single_output)

        with gr.Tab("Generate from Seed"):
            seed_input = gr.Image(label="Seed Image")
            seed_button = gr.Button("Generate from Seed")
            seed_output = gr.Image()
            seed_button.click(generate_from_seed, inputs=seed_input, outputs=seed_output)


    demo.launch()


if __name__ == "__main__":
    main()