app.py

import torch
import torch.nn as nn
import torch.nn.functional as F

import gradio as gr
from PIL import Image
from torchvision import transforms

from huggingface_hub import hf_hub_download

# ============================================================
# CONFIG
# ============================================================

MODEL_REPO = "ASomeoneWhoInterestedWithAI/LookThem_V8-ImageNet100"
MODEL_FILE = "LookThem_V8_Stabilized.pth"

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

# ============================================================
# IMAGENET100 CLASSES
# ============================================================

CLASSES = [

    "bonnet, poke bonnet",
    "green mamba",
    "langur",
    "Doberman, Doberman pinscher",
    "gyromitra",
    "Saluki, gazelle hound",
    "vacuum, vacuum cleaner",
    "window screen",
    "cocktail shaker",
    "garden spider, Aranea diademata",
    "garter snake, grass snake",
    "carbonara",
    "pineapple, ananas",
    "computer keyboard, keypad",
    "tripod",
    "komondor",
    "American lobster, Northern lobster, Maine lobster, Homarus americanus",
    "bannister, banister, balustrade, balusters, handrail",
    "honeycomb",
    "tile roof",
    "papillon",
    "boathouse",
    "stinkhorn, carrion fungus",
    "jean, blue jean, denim",
    "Chihuahua",
    "Chesapeake Bay retriever",
    "robin, American robin, Turdus migratorius",
    "tub, vat",
    "Great Dane",
    "rotisserie",
    "bottlecap",
    "throne",
    "little blue heron, Egretta caerulea",
    "rock crab, Cancer irroratus",
    "Rottweiler",
    "lorikeet",
    "Gila monster, Heloderma suspectum",
    "head cabbage",
    "car wheel",
    "coyote, prairie wolf, brush wolf, Canis latrans",
    "moped",
    "milk can",
    "mixing bowl",
    "toy terrier",
    "chocolate sauce, chocolate syrup",
    "rocking chair, rocker",
    "wing",
    "park bench",
    "ambulance",
    "football helmet",
    "leafhopper",
    "cauliflower",
    "pirate, pirate ship",
    "purse",
    "hare",
    "lampshade, lamp shade",
    "fiddler crab",
    "standard poodle",
    "Shih-Tzu",
    "pedestal, plinth, footstall",
    "gibbon, Hylobates lar",
    "safety pin",
    "English foxhound",
    "chime, bell, gong",
    "American Staffordshire terrier, Staffordshire terrier, American pit bull terrier, pit bull terrier",
    "bassinet",
    "wild boar, boar, Sus scrofa",
    "theater curtain, theatre curtain",
    "dung beetle",
    "hognose snake, puff adder, sand viper",
    "Mexican hairless",
    "mortarboard",
    "Walker hound, Walker foxhound",
    "red fox, Vulpes vulpes",
    "modem",
    "slide rule, slipstick",
    "walking stick, walkingstick, stick insect",
    "cinema, movie theater, movie theatre, movie house, picture palace",
    "meerkat, mierkat",
    "kuvasz",
    "obelisk",
    "harmonica, mouth organ, harp, mouth harp",
    "sarong",
    "mousetrap",
    "hard disc, hard disk, fixed disk",
    "American coot, marsh hen, mud hen, water hen, Fulica americana",
    "reel",
    "pickup, pickup truck",
    "iron, smoothing iron",
    "tabby, tabby cat",
    "ski mask",
    "vizsla, Hungarian pointer",
    "laptop, laptop computer",
    "stretcher",
    "Dutch oven",
    "African hunting dog, hyena dog, Cape hunting dog, Lycaon pictus",
    "boxer",
    "gasmask, respirator, gas helmet",
    "goose",
    "borzoi, Russian wolfhound",
]


class LookThemLayer(nn.Module):
    def __init__(self, num_tokens, in_features, hidden_dim):
        super().__init__()
        self.num_tokens = num_tokens
        self.mod1_w1 = nn.Parameter(torch.randn(num_tokens, in_features, hidden_dim))
        self.mod1_b1 = nn.Parameter(torch.zeros(num_tokens, hidden_dim))
        self.mod1_w2 = nn.Parameter(torch.randn(num_tokens, hidden_dim, 1))
        self.mod1_b2 = nn.Parameter(torch.zeros(num_tokens, 1))
        self.mod2_w1 = nn.Parameter(torch.randn(num_tokens, in_features, hidden_dim))
        self.mod2_b1 = nn.Parameter(torch.zeros(num_tokens, hidden_dim))
        self.mod2_w2 = nn.Parameter(torch.randn(num_tokens, hidden_dim, 1))
        self.mod2_b2 = nn.Parameter(torch.zeros(num_tokens, 1))
        self.trans_w = nn.Parameter(torch.randn(num_tokens, 1, 1))
        self.trans_b = nn.Parameter(torch.zeros(num_tokens, 1))
        self._init_weights()

    def _init_weights(self):
        for w in [self.mod1_w1, self.mod2_w1, self.mod1_w2, self.mod2_w2]:
            nn.init.xavier_uniform_(w) # Better for Tanh/Gelu flow

    def forward(self, x):
        N = self.num_tokens
        h1 = torch.einsum("bti,tij->btj", x, self.mod1_w1) + self.mod1_b1
        out_m1 = torch.einsum("btj,tjk->btk", F.gelu(h1), self.mod1_w2) + self.mod1_b2
        h2 = torch.einsum("bti,tij->btj", x, self.mod2_w1) + self.mod2_b1
        out_m2 = torch.einsum("btj,tjk->btk", F.gelu(h2), self.mod2_w2) + self.mod2_b2

        # Stabilized division
        out_m2_safe = torch.sign(out_m2) * torch.clamp(torch.abs(out_m2), min=1e-6)
        compare = torch.tanh(out_m1.unsqueeze(2) / out_m2_safe.unsqueeze(1))
        compare2 = torch.tanh(out_m1.unsqueeze(1) / out_m2_safe.unsqueeze(2))

        trans_compare = torch.einsum("bije,jef->bijf", compare, self.trans_w) + self.trans_b.view(1, 1, N, 1)
        trans_compare2 = torch.einsum("bije,jef->bijf", compare2, self.trans_w) + self.trans_b.view(1, 1, N, 1)

        interaksi = (trans_compare * x.unsqueeze(2) + trans_compare2 * x.unsqueeze(1)) / 2
        mask = (1.0 - torch.eye(N, device=x.device)).view(1, N, N, 1)
        return (interaksi * mask).sum(dim=2) / (N - 1.0)

class LiteResidualBlock(nn.Module):
    def __init__(self, dim, dropout=0.05):
        super().__init__()
        self.block = nn.Sequential(nn.Linear(dim, dim), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim, dim))
        self.norm = nn.LayerNorm(dim)
    def forward(self, x):
        return self.norm(x + self.block(x))

class LookThemV8(nn.Module):
    def __init__(self):
        super().__init__()
        self.stream_a = nn.Sequential(nn.Conv2d(3, 16, 3, 2, 1), nn.BatchNorm2d(16), nn.GELU(), nn.Conv2d(16, 32, 3, 2, 1), nn.BatchNorm2d(32), nn.GELU(), nn.Conv2d(32, 64, 3, 2, 1), nn.BatchNorm2d(64), nn.GELU(), nn.AdaptiveMaxPool2d((8, 8)))
        self.stream_b = nn.Sequential(nn.Conv2d(3, 16, 3, 2, 1), nn.BatchNorm2d(16), nn.GELU(), nn.Conv2d(16, 32, 3, 1, 1), nn.BatchNorm2d(32), nn.GELU(), nn.Conv2d(32, 64, 3, 1, 1), nn.BatchNorm2d(64), nn.GELU(), nn.AdaptiveMaxPool2d((8, 8)))

        self.lookthemA = LookThemLayer(64, 64, 32)
        self.lookthemB = LookThemLayer(64, 64, 32)
        self.lookthem_comb = LookThemLayer(64, 128, 32)
        self.comb_norm = nn.LayerNorm(128) # Reset distribution after combined LookThem

        self.FFN1 = nn.Conv1d(128, 64, 1)
        self.lookthem2 = LookThemLayer(64, 64, 32)
        self.FFN2 = nn.Conv1d(64, 64, 1)

        self.compressor = nn.Conv1d(64, 16, 1)
        self.input_proj = nn.Linear(64 * 16, 256)
        self.res_blocks = nn.Sequential(LiteResidualBlock(256), LiteResidualBlock(256))
        self.head = nn.Sequential(nn.Linear(256, 128), nn.GELU(), nn.Linear(128, 100))

    def forward(self, x):
        b = x.size(0)
        fa = self.lookthemA(self.stream_a(x).view(b, 64, 64).transpose(1, 2))
        fb = self.lookthemB(self.stream_b(x).view(b, 64, 64).transpose(1, 2))
        x = self.comb_norm(self.lookthem_comb(torch.cat([fa, fb], dim=2)))
        x = x.transpose(1, 2)
        x = self.FFN1(x).transpose(1, 2)
        res = x
        x = self.lookthem2(x).transpose(1, 2)
        x = self.FFN2(x) + res.transpose(1, 2) # Strong residual
        x = self.compressor(x).flatten(1)
        x = self.res_blocks(self.input_proj(x))
        return self.head(x)


# ============================================================
# LOAD MODEL
# ============================================================

print("🧠 Downloading model...")

model_path = hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE
)

print("🧠 Loading model...")

model = LookThemV8().to(DEVICE)

state_dict = torch.load(
    model_path,
    map_location=DEVICE
)

model.load_state_dict(state_dict)

model.eval()

print("✅ Model loaded!")

# ============================================================
# TRANSFORM
# ============================================================

transform = transforms.Compose([

    transforms.Resize((224, 224)),

    transforms.ToTensor(),

    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

# ============================================================
# INFERENCE
# ============================================================

def predict(image):

    image = image.convert("RGB")

    x = transform(image).unsqueeze(0).to(DEVICE)

    with torch.no_grad():

        output = model(x)

        probs = F.softmax(output, dim=1)

        top_probs, top_idx = torch.topk(probs, 5)

    results = {}

    for p, idx in zip(top_probs[0], top_idx[0]):

        results[CLASSES[idx.item()]] = float(p.item())

    return results

# ============================================================
# GRADIO UI
# ============================================================

demo = gr.Interface(

    fn=predict,

    inputs=gr.Image(type="pil"),

    outputs=gr.Label(num_top_classes=5),

    title="LookThem V8 ImageNet100",

    description="Tiny relational vision model using LookThem architecture "
)

demo.launch()