app.py

import os
import torch
import torch.nn as nn
from PIL import Image
import torchvision.transforms as transforms
import numpy as np
import gradio as gr

# ── Professional "Cyber-Dark" CSS ─────────────────────────────
custom_css = """
.gradio-container {background-color: #050505;}
.feedback-card {border: 1px solid #6366f1; padding: 20px; border-radius: 12px; background: #0f172a; margin: 10px 0;}
#header-text {text-align: center; background: linear-gradient(to right, #818cf8, #c084fc); -webkit-background-clip: text; -webkit-text-fill-color: transparent;}
button.primary {background: linear-gradient(90deg, #6366f1, #a855f7) !important; color: white !important; font-weight: bold !important; border: none !important;}
.tabs {border: none !important;}
footer {visibility: hidden}
"""

# [Model classes: PatchifyAndMask, TransformerBlock, ViT_Encoder, ViT_Decoder remain identical]
class PatchifyAndMask(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3):
        super().__init__()
        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = (img_size // patch_size) ** 2
        self.patch_dim = in_channels * patch_size * patch_size

    def patchify(self, imgs):
        p = self.patch_size
        h = w = self.img_size // p
        x = imgs.reshape(shape=(imgs.shape[0], 3, h, p, w, p))
        x = torch.einsum('nchpwq->nhwpqc', x)
        x = x.reshape(shape=(imgs.shape[0], h * w, p**2 * 3))
        return x

    def random_masking(self, x, mask_ratio=0.75):
        N, L, D = x.shape
        len_keep = int(L * (1 - mask_ratio))
        noise = torch.rand(N, L, device=x.device)
        ids_shuffle = torch.argsort(noise, dim=1)
        ids_restore = torch.argsort(ids_shuffle, dim=1)
        ids_keep = ids_shuffle[:, :len_keep]
        x_kept = torch.gather(x, dim=1, index=ids_keep.unsqueeze(-1).repeat(1, 1, D))
        mask = torch.ones([N, L], device=x.device)
        mask[:, :len_keep] = 0
        mask = torch.gather(mask, dim=1, index=ids_restore)
        return x_kept, mask, ids_restore

class TransformerBlock(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super().__init__()
        self.norm1 = nn.LayerNorm(embed_dim)
        self.norm2 = nn.LayerNorm(embed_dim)
        self.attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
        self.mlp = nn.Sequential(
            nn.Linear(embed_dim, embed_dim * 4),
            nn.GELU(),
            nn.Linear(embed_dim * 4, embed_dim)
        )

    def forward(self, x):
        attn_output, _ = self.attn(self.norm1(x), self.norm1(x), self.norm1(x))
        x = x + attn_output
        x = x + self.mlp(self.norm2(x))
        return x

class ViT_Encoder(nn.Module):
    def __init__(self, patch_dim=768, embed_dim=768, depth=12, num_heads=12, num_patches=196):
        super().__init__()
        self.patch_embed = nn.Linear(patch_dim, embed_dim)
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim))
        self.blocks = nn.ModuleList([
            TransformerBlock(embed_dim, num_heads) for _ in range(depth)
        ])
        self.norm = nn.LayerNorm(embed_dim)

    def forward(self, x, patcher_module):
        x = self.patch_embed(x)
        x = x + self.pos_embed
        x_visible, mask, ids_restore = patcher_module.random_masking(x, mask_ratio=0.75)
        for block in self.blocks:
            x_visible = block(x_visible)
        x_visible = self.norm(x_visible)
        return x_visible, mask, ids_restore

class ViT_Decoder(nn.Module):
    def __init__(self, encoder_dim=768, decoder_dim=384, depth=12, num_heads=6, num_patches=196, patch_dim=768):
        super().__init__()
        self.decoder_embed = nn.Linear(encoder_dim, decoder_dim)
        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim))
        self.decoder_pos_embed = nn.Parameter(torch.zeros(1, num_patches, decoder_dim))
        self.blocks = nn.ModuleList([
            TransformerBlock(embed_dim=decoder_dim, num_heads=num_heads) for _ in range(depth)
        ])
        self.norm = nn.LayerNorm(decoder_dim)
        self.pred = nn.Linear(decoder_dim, patch_dim)

    def forward(self, x, ids_restore):
        x = self.decoder_embed(x)
        B = x.shape[0]
        num_visible = x.shape[1]
        num_total = ids_restore.shape[1]
        num_masks = num_total - num_visible
        mask_tokens = self.mask_token.repeat(B, num_masks, 1)
        x_full = torch.cat([x, mask_tokens], dim=1)
        ids_restore_expanded = ids_restore.unsqueeze(-1).repeat(1, 1, x_full.shape[2])
        x_unshuffled = torch.gather(x_full, dim=1, index=ids_restore_expanded)
        x_unshuffled = x_unshuffled + self.decoder_pos_embed
        for block in self.blocks:
            x_unshuffled = block(x_unshuffled)
        x_unshuffled = self.norm(x_unshuffled)
        predictions = self.pred(x_unshuffled)
        return predictions

def unpatchify(x, patch_size=16, img_size=224):
    p = patch_size
    h = w = img_size // p
    x = x.reshape(shape=(x.shape[0], h, w, p, p, 3))
    x = torch.einsum('nhwpqc->nchpwq', x)
    imgs = x.reshape(shape=(x.shape[0], 3, h * p, w * p))
    return imgs

class MaskedAutoencoderApp(nn.Module):
    def __init__(self):
        super().__init__()
        self.patcher = PatchifyAndMask(img_size=224, patch_size=16, in_channels=3)
        self.encoder = ViT_Encoder(patch_dim=768, embed_dim=768, depth=12, num_heads=12)
        self.decoder = ViT_Decoder(encoder_dim=768, decoder_dim=384, depth=12, num_heads=6)

    def forward(self, imgs, mask_ratio):
        target_patches = self.patcher.patchify(imgs)
        x = self.encoder.patch_embed(target_patches)
        x = x + self.encoder.pos_embed
        x_visible, mask, ids_restore = self.patcher.random_masking(x, mask_ratio=mask_ratio)
        for block in self.encoder.blocks:
            x_visible = block(x_visible)
        x_visible = self.encoder.norm(x_visible)
        predictions = self.decoder(x_visible, ids_restore)
        return predictions, target_patches, mask

# ── Load Model ────────────────────────────────────────────────
model = MaskedAutoencoderApp()
try:
    model.load_state_dict(torch.load("best_mae_weights.pth", map_location=torch.device("cpu")))
except Exception as e:
    print(f"Error loading weights: {e}")
model.eval()

# ── Inference ─────────────────────────────────────────────────
def process_image(input_image, mask_ratio_percent):
    if input_image is None: return None, None, None
    mask_ratio = mask_ratio_percent / 100.0
    image = input_image.convert('RGB')
    transform = transforms.Compose([transforms.Resize((224, 224)), transforms.ToTensor()])
    input_tensor = transform(image).unsqueeze(0)
    
    with torch.no_grad():
        preds, targets, mask = model(input_tensor, mask_ratio=mask_ratio)
        mask_expanded = mask.unsqueeze(-1).repeat(1, 1, targets.shape[2])
        masked_input = targets * (1 - mask_expanded)
        masked_img_np = unpatchify(masked_input).squeeze().numpy()
        final_recon = preds * mask_expanded + targets * (1 - mask_expanded)
        recon_img_np = unpatchify(final_recon).squeeze().numpy()
        orig_img_np = unpatchify(targets).squeeze().numpy()

    return (np.clip(np.transpose(masked_img_np, (1, 2, 0)), 0, 1),
            np.clip(np.transpose(recon_img_np, (1, 2, 0)), 0, 1),
            np.clip(np.transpose(orig_img_np, (1, 2, 0)), 0, 1))

# ── Enhanced UI Construction ──────────────────────────────────
with gr.Blocks(css=custom_css, theme=gr.themes.Monochrome()) as demo:
    with gr.Column():
        gr.Markdown("# 💎 PIXEL REVIVE: AI RECONSTRUCTION", elem_id="header-text")
        gr.Markdown("Self-supervised Masked Autoencoder (MAE) designed to reconstruct images by learning global context from visible patches.")

    with gr.Tabs(elem_classes="tabs"):
        with gr.TabItem("🚀 Reconstruction Studio"):
            with gr.Row():
                with gr.Column(scale=1, variant="panel"):
                    gr.Markdown("### Processing Settings")
                    img_input = gr.Image(type="pil", label="Source Image")
                    mask_slider = gr.Slider(10, 90, value=75, step=5, label="Masking Intensity (%)")
                    submit_btn = gr.Button("✨ RECONSTRUCT", variant="primary")
                    gr.Markdown("*(MAE models are often pre-trained with 75% masking)*")
                
                with gr.Column(scale=2):
                    with gr.Row():
                        out_masked = gr.Image(label="Masked Input")
                        out_recon = gr.Image(label="AI Output")
                    out_orig = gr.Image(label="Reference Image", height=200)

        with gr.TabItem("📊 Technical Architecture"):
            with gr.Row():
                with gr.Column(elem_classes="feedback-card"):
                    gr.Markdown("### Vision Transformer (ViT) Encoder")
                    gr.Markdown("Only processes visible patches. By masking 75% of the image, the encoder is forced to learn highly descriptive structural representations.")
                with gr.Column(elem_classes="feedback-card"):
                    gr.Markdown("### Lightweight Decoder")
                    gr.Markdown("Reconstructs the original pixels by combining encoded visible patches with learnable 'mask tokens' to fill in the gaps.")

    submit_btn.click(process_image, [img_input, mask_slider], [out_masked, out_recon, out_orig])

if __name__ == "__main__":
    demo.launch()