# import spaces

import os
import math
import random

import gradio as gr
import torch
from torch import nn
from torch.nn import functional as F
from PIL import Image, ImageDraw, ImageFont
import torchvision.transforms as transforms
from torch.amp import autocast
from diffusers import AutoencoderKL
from huggingface_hub import hf_hub_download

from indic_transliteration import sanscript
from indic_transliteration.sanscript import transliterate

from torchvision.models import mobilenet_v2, MobileNet_V2_Weights

# CONFIG
CKPT_REPO_ID = "keysun89/HW_Hindi_Model"     
CKPT_FILENAME = "ckpt_epoch_434.pt"
VAE_PATH = "runwayml/stable-diffusion-v1-5"                              
FONT_PATH = "NotoSansDevanagari-Regular.ttf" 

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

TIMESTEPS = 1000
betas = torch.linspace(1e-4, 0.02, TIMESTEPS, device=device)
alphas = 1.0 - betas
alpha_bars = torch.cumprod(alphas, dim=0)


class Resblock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        groups_in = min(32, in_channels) if in_channels >= 8 else in_channels
        groups_out = min(32, out_channels) if out_channels >= 8 else out_channels
        self.groupnorm_1 = nn.GroupNorm(groups_in, in_channels)
        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.groupnorm_2 = nn.GroupNorm(groups_out, out_channels)
        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.residual_layer = (nn.Identity() if in_channels == out_channels
                                else nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0))

    def forward(self, x):
        residual = x
        x = self.groupnorm_1(x); x = F.silu(x); x = self.conv1(x)
        x = self.groupnorm_2(x); x = F.silu(x); x = self.conv2(x)
        x = self.conv3(x)
        return x + self.residual_layer(residual)


class SelfAttention(nn.Module):
    def __init__(self, channels, n_heads=8):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = channels // n_heads
        self.qkv = nn.Linear(channels, channels * 3)
        self.proj = nn.Linear(channels, channels)

    def forward(self, x):
        b, c, h, w = x.shape
        x_flat = x.view(b, c, h * w).transpose(1, 2)
        qkv = self.qkv(x_flat)
        q, k, v = qkv.chunk(3, dim=-1)
        q = q.view(b, -1, self.n_heads, self.d_head).transpose(1, 2)
        k = k.view(b, -1, self.n_heads, self.d_head).transpose(1, 2)
        v = v.view(b, -1, self.n_heads, self.d_head).transpose(1, 2)
        attn = torch.softmax(q @ k.transpose(-1, -2) / math.sqrt(self.d_head), dim=-1)
        out = attn @ v
        out = out.transpose(1, 2).reshape(b, h * w, c)
        out = self.proj(out)
        return out.transpose(1, 2).reshape(b, c, h, w)


class SelfAttentionBlock(nn.Module):
    def __init__(self, channels):
        super().__init__()
        groups = min(32, channels) if channels >= 8 else channels
        self.norm = nn.GroupNorm(groups, channels)
        self.attn = SelfAttention(channels)

    def forward(self, x):
        return x + self.attn(self.norm(x))


class StyleProxyHead(nn.Module):
    """Kept only for checkpoint state_dict compatibility (strict=False loading);
    not used in the inference forward path."""
    def __init__(self, in_channels=64, embed_dim=512, mode="column", mask_ratio=0.5):
        super().__init__()
        self.mode = mode
        self.mask_ratio = mask_ratio
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(in_channels, embed_dim)

    def forward(self, feat_map):
        x = self.pool(feat_map).flatten(1)
        x = self.fc(x)
        return F.normalize(x, p=2, dim=1)


class Ver_Style(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(512, 256, kernel_size=3, padding=1)
        self.res1 = Resblock(256, 256)
        self.res2 = Resblock(256, 64)
        self.attn1 = SelfAttentionBlock(64)
        self.attn2 = SelfAttentionBlock(64)
        self.attn3 = SelfAttentionBlock(64)
        self.res3 = Resblock(64, 64)

    def forward(self, x):
        x = self.conv1(x); x = self.res1(x); x = self.res2(x)
        x = self.attn1(x); x = self.attn2(x); x = self.attn3(x)
        return self.res3(x)


class Hor_Style(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(512, 256, kernel_size=3, padding=1)
        self.res1 = Resblock(256, 256)
        self.res2 = Resblock(256, 64)
        self.attn1 = SelfAttentionBlock(64)
        self.attn2 = SelfAttentionBlock(64)
        self.attn3 = SelfAttentionBlock(64)
        self.res3 = Resblock(64, 64)

    def forward(self, x):
        x = self.conv1(x); x = self.res1(x); x = self.res2(x)
        x = self.attn1(x); x = self.attn2(x); x = self.attn3(x)
        return self.res3(x)


class MobileNetStride8Backbone(nn.Module):
    def __init__(self, out_channels=512):
        super().__init__()
        mnet = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
        self.features = mnet.features[:7]
        self.project = nn.Conv2d(32, out_channels, kernel_size=1)

    def forward(self, x):
        return self.project(self.features(x))


class StyleEncoder(nn.Module):
    def __init__(self, embed_dim=512, mask_ratio=0.5):
        super().__init__()
        self.backbone = MobileNetStride8Backbone(out_channels=512)
        self.ver = Ver_Style()
        self.hor = Hor_Style()
        self.ver_proxy_head = StyleProxyHead(64, embed_dim, mode="column", mask_ratio=mask_ratio)
        self.hor_proxy_head = StyleProxyHead(64, embed_dim, mode="row", mask_ratio=mask_ratio)
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(512, embed_dim)

    def forward(self, x):
        feat = self.backbone(x)
        ver_map = self.ver(feat)
        hor_map = self.hor(feat)
        ver_emb = self.ver_proxy_head(ver_map)
        hor_emb = self.hor_proxy_head(hor_map)
        global_emb = self.pool(feat).flatten(1)
        global_emb = F.normalize(self.fc(global_emb), p=2, dim=1)
        return ver_map, hor_map, ver_emb, hor_emb, global_emb


class UnifontTextEncoder(nn.Module):
    def __init__(self, font_path=FONT_PATH, image_size=(64, 1024)):
        super().__init__()
        self.font_path = font_path
        self.image_size = image_size
        self.transform = transforms.Compose([
            transforms.Resize(image_size),
            transforms.Grayscale(num_output_channels=3),
            transforms.ToTensor(),
            transforms.Normalize([0.5] * 3, [0.5] * 3),
        ])
        try:
            self.font = ImageFont.truetype(font_path, 28)
        except Exception:
            self.font = ImageFont.load_default()

    def render_text(self, text):
        H, W = self.image_size
        img = Image.new('RGB', (W, H), (255, 255, 255))
        draw = ImageDraw.Draw(img)
        font_size, margin, min_font_size = 28, 20, 8
        font = self._load_font(font_size)
        while font_size > min_font_size:
            bbox = draw.textbbox((0, 0), text, font=font)
            if (bbox[2] - bbox[0]) <= (W - margin):
                break
            font_size -= 2
            font = self._load_font(font_size)
        bbox = draw.textbbox((0, 0), text, font=font)
        text_h, text_w = bbox[3] - bbox[1], bbox[2] - bbox[0]
        x = max(0, (W - text_w) // 2 - bbox[0])
        y = max(0, (H - text_h) // 2 - bbox[1])
        draw.text((x, y), text, font=font, fill=(0, 0, 0))
        return img

    def _load_font(self, font_size):
        try:
            return ImageFont.truetype(self.font_path, font_size, layout_engine=ImageFont.Layout.RAQM)
        except Exception:
            try:
                return ImageFont.truetype(self.font_path, font_size)
            except Exception:
                return self.font

    def forward(self, texts, device=None):
        imgs = torch.stack([self.transform(self.render_text(t)) for t in texts])
        return imgs.to(device) if device is not None else imgs


class ContentEncoder(nn.Module):
    def __init__(self, font_path=FONT_PATH):
        super().__init__()
        self.unifont = UnifontTextEncoder(font_path)
        mnet = mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
        self.backbone = mnet.features[:7]
        self.proj = nn.Conv2d(32, 64, kernel_size=1)
        self.attn_head = nn.Sequential(
            SelfAttentionBlock(64), SelfAttentionBlock(64), SelfAttentionBlock(64),
            Resblock(64, 64),
        )
        self.pos_embed = nn.Parameter(torch.zeros(1, 64, 1, 128))

    def forward(self, text):
        dev = next(self.parameters()).device
        x = self.unifont(text, device=dev)
        x = self.backbone(x)
        x = self.proj(x)
        x = x + self.pos_embed
        return self.attn_head(x)


class CrossAttention_unet(nn.Module):
    def __init__(self, ch_1, ch_2=64):
        super().__init__()
        self.ch_1 = ch_1
        self.q = nn.Linear(ch_1, ch_1)
        self.k = nn.Linear(ch_2, ch_1)
        self.v = nn.Linear(ch_2, ch_1)
        self.proj = nn.Linear(ch_1, ch_1)

    def forward(self, x, cond):
        is_4d = x.dim() == 4
        if is_4d:
            b, c, h, w = x.shape
            x = x.view(b, c, h * w).transpose(1, 2)
        if cond.dim() == 4:
            cond = cond.flatten(2).transpose(1, 2)
        q, k, v = self.q(x), self.k(cond), self.v(cond)
        attn = torch.softmax(q @ k.transpose(1, 2) / math.sqrt(self.ch_1), dim=-1)
        out = self.proj(attn @ v)
        if is_4d:
            out = out.transpose(1, 2).view(b, c, h, w)
        return out


class Blender(nn.Module):
    def __init__(self):
        super().__init__()
        self.cross_ver = CrossAttention_unet(64, 64)
        self.self_ver_1 = SelfAttentionBlock(64)
        self.self_ver_2 = SelfAttentionBlock(64)
        self.cross_hor = CrossAttention_unet(64, 64)
        self.self_hor_1 = SelfAttentionBlock(64)
        self.self_hor_2 = SelfAttentionBlock(64)
        self.res_out = Resblock(64, 64)

    def forward(self, Q, S_ver, S_hor):
        cond = self.cross_ver(Q, S_ver)
        cond = self.self_ver_1(cond)
        cond = self.self_ver_2(cond)
        cond = self.cross_hor(cond, S_hor)
        cond = self.self_hor_1(cond)
        cond = self.self_hor_2(cond)
        cond = cond + Q * 0.5
        return self.res_out(cond)


class TimeEmbedding(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        half = self.dim // 2
        freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / (half - 1))
        args = t[:, None] * freqs[None, :]
        return torch.cat([torch.sin(args), torch.cos(args)], dim=1)


class UNET_ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, time_dim=512):
        super().__init__()
        self.groupnorm_feature = nn.GroupNorm(32, in_channels)
        self.conv_feature = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
        self.groupnorm_merged = nn.GroupNorm(32, out_channels)
        self.conv_merged = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
        self.time_proj = nn.Linear(time_dim, out_channels)
        self.residual_layer = (nn.Identity() if in_channels == out_channels
                                else nn.Conv2d(in_channels, out_channels, kernel_size=1, padding=0))

    def forward(self, feature, time):
        residue = feature
        feature = self.groupnorm_feature(feature); feature = F.silu(feature); feature = self.conv_feature(feature)
        time = F.silu(time); time = self.time_proj(time)
        merged = feature + time.unsqueeze(-1).unsqueeze(-1)
        merged = self.groupnorm_merged(merged); merged = F.silu(merged); merged = self.conv_merged(merged)
        return merged + self.residual_layer(residue)


class UNET_AttentionBlock(nn.Module):
    def __init__(self, channels, cond_dim=64, geglu_mult=2):
        super().__init__()
        self.conv_input = nn.Conv2d(channels, channels, kernel_size=1, padding=0)
        self.layernorm_1 = nn.LayerNorm(channels)
        self.attention_1 = SelfAttention(channels)
        self.layernorm_2 = nn.LayerNorm(channels)
        self.attention_2 = CrossAttention_unet(channels, cond_dim)
        self.layernorm_3 = nn.LayerNorm(channels)
        hidden = geglu_mult * channels
        self.linear_geglu_1 = nn.Linear(channels, hidden * 2)
        self.linear_geglu_2 = nn.Linear(hidden, channels)
        self.conv_output = nn.Conv2d(channels, channels, kernel_size=1, padding=0)

    def forward(self, x, cond):
        residue_long = x
        x = self.conv_input(x)
        b, c, h, w = x.shape
        x = x.view((b, c, h * w)).transpose(1, 2)
        residue_short = x
        x = self.layernorm_1(x)
        x_spatial = x.transpose(1, 2).reshape(b, c, h, w)
        x_spatial = self.attention_1(x_spatial)
        x = x_spatial.reshape(b, c, h * w).transpose(1, 2)
        x += residue_short
        residue_short = x
        x = self.layernorm_2(x)
        x = self.attention_2(x, cond)
        x += residue_short
        residue_short = x
        x = self.layernorm_3(x)
        x, gate = self.linear_geglu_1(x).chunk(2, dim=-1)
        x = x * F.gelu(gate)
        x = self.linear_geglu_2(x)
        x += residue_short
        x = x.transpose(1, 2).view((b, c, h, w))
        return self.conv_output(x) + residue_long


class SwitchSequential(nn.Sequential):
    def forward(self, x, cond, time):
        for layer in self:
            if isinstance(layer, UNET_AttentionBlock):
                x = layer(x, cond)
            elif isinstance(layer, UNET_ResidualBlock):
                x = layer(x, time)
            else:
                x = layer(x)
        return x


class width_Upsample(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
    def forward(self, x):
        return self.conv(F.interpolate(x, scale_factor=(1.0, 2.0), mode='nearest'))

class width_Downsample(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Conv2d(channels, channels, kernel_size=3, padding=1, stride=(1, 2))
    def forward(self, x):
        return self.conv(x)

class Height_Upsample(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
    def forward(self, x):
        return self.conv(F.interpolate(x, scale_factor=(2.0, 1.0), mode='nearest'))

class Height_Downsample(nn.Module):
    def __init__(self, channels):
        super().__init__()
        self.conv = nn.Conv2d(channels, channels, kernel_size=3, padding=1, stride=(2, 1))
    def forward(self, x):
        return self.conv(x)


class Unet(nn.Module):
    def __init__(self, c1=224, c2=448, c3=896, time_dim=896, cond_dim=64, geglu_mult=2):
        super().__init__()
        def RB(i, o): return UNET_ResidualBlock(i, o, time_dim=time_dim)
        def AB(c): return UNET_AttentionBlock(c, cond_dim, geglu_mult=geglu_mult)

        self.encoders = nn.ModuleList([
            SwitchSequential(nn.Conv2d(4, c1, kernel_size=3, padding=1)),
            SwitchSequential(RB(c1, c1), AB(c1)),
            SwitchSequential(width_Downsample(c1)),
            SwitchSequential(RB(c1, c2), AB(c2)),
            SwitchSequential(width_Downsample(c2)),
            SwitchSequential(RB(c2, c3), AB(c3)),
            SwitchSequential(width_Downsample(c3)),
            SwitchSequential(Height_Downsample(c3)),
            SwitchSequential(RB(c3, c3), AB(c3)),
        ])
        self.bottleneck = SwitchSequential(RB(c3, c3), AB(c3), RB(c3, c3))
        self.decoders = nn.ModuleList([
            SwitchSequential(RB(c3 * 2, c3), AB(c3)),
            SwitchSequential(RB(c3 * 2, c3), Height_Upsample(c3)),
            SwitchSequential(RB(c3 * 2, c3), AB(c3), width_Upsample(c3)),
            SwitchSequential(RB(c3 * 2, c2), AB(c2)),
            SwitchSequential(RB(c2 * 2, c2), width_Upsample(c2)),
            SwitchSequential(RB(c2 * 2, c1), AB(c1)),
            SwitchSequential(RB(c1 * 2, c1), width_Upsample(c1)),
            SwitchSequential(RB(c1 * 2, c1), AB(c1)),
            SwitchSequential(RB(c1 * 2, c1)),
        ])

    def forward(self, x, cond, time):
        skip_connections = []
        for layers in self.encoders:
            x = layers(x, cond, time)
            skip_connections.append(x)
        x = self.bottleneck(x, cond, time)
        for layers in self.decoders:
            x = torch.cat((x, skip_connections.pop()), dim=1)
            x = layers(x, cond, time)
        return x


class GlobalEmbeddingHead(nn.Module):
    def __init__(self, in_channels=4, embed_dim=512):
        super().__init__()
        self.pool = nn.AdaptiveAvgPool2d((1, 1))
        self.fc = nn.Linear(in_channels, embed_dim)

    def forward(self, x):
        x = self.pool(x).flatten(1)
        return F.normalize(self.fc(x), p=2, dim=1)


class UNET_OutputLayer(nn.Module):
    def __init__(self, in_channels, out_channels):
        super().__init__()
        self.groupnorm = nn.GroupNorm(32, in_channels)
        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)

    def forward(self, x):
        x = self.groupnorm(x); x = F.silu(x)
        return self.conv(x)


class Diffusion(nn.Module):
    def __init__(self, font_path=FONT_PATH):
        super().__init__()
        self.style_encoder = StyleEncoder()
        self.content_encoder = ContentEncoder(font_path)
        self.blender = Blender()
        self.time_embedding = TimeEmbedding(896)
        self.unet = Unet()
        self.final = UNET_OutputLayer(224, 4)
        self.embedding_head = GlobalEmbeddingHead(in_channels=4, embed_dim=512)


# LOAD MODEL + VAE (runs once at Space startup)

def _load_partial(module, state_dict, name):
    result = module.load_state_dict(state_dict, strict=False)
    if result.missing_keys:
        print(f"[{name}] missing keys: {result.missing_keys}")
    if result.unexpected_keys:
        print(f"[{name}] unexpected keys: {result.unexpected_keys}")


print("Downloading checkpoint from HF Hub...")
ckpt_path = hf_hub_download(repo_id=CKPT_REPO_ID, filename=CKPT_FILENAME)

print("Loading VAE...")
vae = AutoencoderKL.from_pretrained(VAE_PATH, subfolder="vae").to(device)
vae.eval()
for p in vae.parameters():
    p.requires_grad = False

print("Loading Diffusion model...")
model = Diffusion(font_path=FONT_PATH).to(device)
ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)

_load_partial(model.style_encoder, ckpt["style_encoder"], "style_encoder")
_load_partial(model.content_encoder, ckpt["content_encoder"], "content_encoder")
_load_partial(model.blender, ckpt["blender"], "blender")
_load_partial(model.unet, ckpt["unet"], "unet")
_load_partial(model.final, ckpt["final"], "final")
if "embedding_head" in ckpt:
    _load_partial(model.embedding_head, ckpt["embedding_head"], "embedding_head")

model.eval()
print(f"[LOADED] checkpoint from epoch {ckpt.get('epoch', '?')}")


# INFERENCE
# @spaces.GPU(duration=120)   # bump duration since up to 1000 steps may need more time
def generate(style_image, text, seed, num_steps, progress=gr.Progress()):
    if style_image is None:
        raise gr.Error("Please upload a style reference image.")
    if not text or not text.strip():
        raise gr.Error("Please enter some text to render.")

    seed = int(seed) if seed not in (None, "") else random.randint(0, 999999)
    torch.manual_seed(seed)
    random.seed(seed)

    display_text = text
    if text.isascii():
        display_text = transliterate(text.lower(), sanscript.ITRANS, sanscript.DEVANAGARI)

    transform = transforms.Compose([
        transforms.Resize((64, 1024)),
        transforms.ToTensor(),
        transforms.Normalize([0.5] * 3, [0.5] * 3),
    ])

    style_img = transform(style_image.convert("RGB")).unsqueeze(0).to(device)

    progress(0, desc="Encoding style and text...")

    with torch.no_grad():
        with autocast(device_type=device.type):
            ver_map, hor_map, _, _, _ = model.style_encoder(style_img)
            Q = model.content_encoder([display_text])
            cond = model.blender(Q, ver_map, hor_map)

        latent = torch.randn(1, 4, 8, 128, device=device)

        NUM_DDIM_STEPS = int(num_steps)
        ddim_timesteps = torch.linspace(TIMESTEPS - 1, 0, NUM_DDIM_STEPS, dtype=torch.long, device=device)

        for i in range(len(ddim_timesteps)):
            t = ddim_timesteps[i].item()
            t_prev = ddim_timesteps[i + 1].item() if i + 1 < len(ddim_timesteps) else -1

            t_tensor = torch.tensor([t], device=device)
            with autocast(device_type=device.type):
                time_emb = model.time_embedding(t_tensor)
                pred_noise = model.unet(latent, cond, time_emb)
                pred_noise = model.final(pred_noise)

            alpha_bar_t = alpha_bars[t]
            alpha_bar_t_prev = alpha_bars[t_prev] if t_prev >= 0 else torch.tensor(1.0, device=device)

            x0_pred = (latent - torch.sqrt(1 - alpha_bar_t) * pred_noise) / torch.sqrt(alpha_bar_t)
            x0_pred = torch.clamp(x0_pred, -3.0, 3.0)

            latent = torch.sqrt(alpha_bar_t_prev) * x0_pred + torch.sqrt(1 - alpha_bar_t_prev) * pred_noise

            progress((i + 1) / NUM_DDIM_STEPS, desc=f"Step {i + 1}/{NUM_DDIM_STEPS}")

        latent = latent / 0.18215
        with autocast(device_type=device.type):
            img = vae.decode(latent).sample

    progress(1.0, desc="Done")

    img = (img.clamp(-1, 1) + 1) / 2
    img = img.squeeze(0).permute(1, 2, 0).cpu().numpy()
    img = (img * 255).astype("uint8")

    return Image.fromarray(img), display_text, seed
    
# GRADIO UI
with gr.Blocks(title="DiffBrush — Hindi Handwriting Generation") as demo:
    gr.Markdown(
        """
        # DiffBrush — Devanagari Handwriting Generation
        Upload a handwriting style reference image and enter text
        (English/ITRANS transliteration or native Devanagari).
        """
    )

    with gr.Row():
        with gr.Column():
            style_input = gr.Image(label="Style reference image", type="pil")
            text_input = gr.Textbox(
                label="Text to render",
                placeholder="Type in English (ITRANS) or Devanagari, e.g. 'namaste' or 'नमस्ते'",
            )
            seed_input = gr.Textbox(label="Seed (optional)", placeholder="Leave blank for random")
            steps_input = gr.Slider(
                minimum=100, maximum=1000, value=100, step=50,
                label="DDIM sampling steps (higher = slower, potentially better quality)",
            )
            submit_btn = gr.Button("Generate", variant="primary")

        with gr.Column():
            output_image = gr.Image(label="Generated handwriting")
            transliterated_text = gr.Textbox(label="Text used (after transliteration)", interactive=False)
            used_seed = gr.Textbox(label="Seed used", interactive=False)

    submit_btn.click(
        fn=generate,
        inputs=[style_input, text_input, seed_input, steps_input],
        outputs=[output_image, transliterated_text, used_seed],
    )

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