Create app.py

app.py

@@ -0,0 +1,231 @@
+"""
+AKASHA - Gradio Demo for Hugging Face Spaces
+https://huggingface.co/spaces/vedaco/AKASHA
+"""
+
+import gradio as gr
+import tensorflow as tf
+import numpy as np
+import os
+import json
+from PIL import Image
+
+os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+
+
+# ── Load model ──────────────────────────────────────────────
+
+def load_model():
+    config_path = "config.json"
+    if os.path.exists(config_path):
+        with open(config_path, "r") as f:
+            config = json.load(f)
+    else:
+        from akasha.utils import create_default_config
+        config = create_default_config()
+
+    from akasha.tokenizer import VQVAE
+    from akasha.model import AKASHAModel
+
+    # Build VQVAE
+    vqvae = VQVAE(config)
+    img_size = config["model"]["tokenizer"]["image_size"]
+    dummy = tf.zeros([1, img_size, img_size, 3])
+    vqvae(dummy)
+
+    vqvae_path = "vqvae.weights.h5"
+    if os.path.exists(vqvae_path):
+        vqvae.load_weights(vqvae_path)
+        print("✅ VQVAE weights loaded")
+    else:
+        print("⚠️  VQVAE weights not found — using random weights (demo mode)")
+
+    # Build Transformer
+    model = AKASHAModel(config)
+    model.set_vqvae(vqvae)
+
+    bos = config["model"]["tokenizer"]["num_tokens"]
+    dummy_tok = tf.cast(tf.fill([1, 10], bos), tf.int32)
+    model.transformer(dummy_tok)
+
+    transformer_path = "transformer.weights.h5"
+    if os.path.exists(transformer_path):
+        model.transformer.load_weights(transformer_path)
+        print("✅ Transformer weights loaded")
+    else:
+        print("⚠️  Transformer weights not found — using random weights (demo mode)")
+
+    return model, config
+
+
+print("🚀 Loading AKASHA model...")
+try:
+    model, config = load_model()
+    MODEL_LOADED = True
+    print("✅ AKASHA ready!")
+except Exception as e:
+    print(f"❌ Failed to load model: {e}")
+    MODEL_LOADED = False
+    model, config = None, None
+
+
+# ── Generation function ─────────────────────────────────────
+
+def generate_images(num_images, temperature, top_k, top_p, seed, progress=gr.Progress()):
+    if not MODEL_LOADED:
+        return None, "❌ Model not loaded. Check logs."
+
+    if seed >= 0:
+        tf.random.set_seed(int(seed))
+        np.random.seed(int(seed))
+
+    num_images = int(num_images)
+    top_k = int(top_k)
+
+    def progress_cb(step, total):
+        progress(step / total, desc=f"Token {step}/{total}")
+
+    try:
+        images, tokens = model.generate(
+            num_images=num_images,
+            temperature=temperature,
+            top_k=top_k,
+            top_p=top_p,
+            callback=progress_cb,
+        )
+
+        pil_images = []
+        for i in range(num_images):
+            img = images[i].numpy()
+            img = (img * 255).clip(0, 255).astype(np.uint8)
+            pil_images.append(Image.fromarray(img))
+
+        info = (
+            f"✅ Generated {num_images} image(s) | "
+            f"temp={temperature} | top_k={top_k} | top_p={top_p}"
+        )
+        if seed >= 0:
+            info += f" | seed={int(seed)}"
+        return pil_images, info
+
+    except Exception as e:
+        return None, f"❌ Generation failed: {str(e)}"
+
+
+# ── Reconstruction function ─────────────────────────────────
+
+def reconstruct_image(input_image, progress=gr.Progress()):
+    if not MODEL_LOADED:
+        return None, "❌ Model not loaded."
+    if input_image is None:
+        return None, "Please upload an image."
+
+    try:
+        img_size = config["model"]["tokenizer"]["image_size"]
+        img = tf.image.resize(input_image, [img_size, img_size])
+        img = tf.cast(img, tf.float32) / 255.0
+        img = tf.expand_dims(img, 0)
+
+        reconstructed, tokens, vq_loss = model.vqvae(img, training=False)
+
+        recon = reconstructed[0].numpy()
+        recon = (recon * 255).clip(0, 255).astype(np.uint8)
+
+        n_tokens = int(tokens.shape[1] * tokens.shape[2])
+        unique = len(np.unique(tokens.numpy()))
+        codebook_size = config["model"]["tokenizer"]["num_tokens"]
+        info = (
+            f"Tokens: {n_tokens} | "
+            f"Unique codes: {unique}/{codebook_size} | "
+            f"VQ Loss: {float(vq_loss):.4f}"
+        )
+        return Image.fromarray(recon), info
+
+    except Exception as e:
+        return None, f"❌ Reconstruction failed: {str(e)}"
+
+
+# ── Build UI ────────────────────────────────────────────────
+
+css = """
+.gradio-container { max-width: 1000px !important; }
+h1 { text-align: center; }
+"""
+
+with gr.Blocks(title="AKASHA — Image Generation", theme=gr.themes.Soft(), css=css) as demo:
+
+    gr.Markdown("""
+    # 🎨 AKASHA — Autoregressive Image Generation
+    Generate images **token-by-token** using a transformer that predicts the next image token.
+
+    Built with **TensorFlow** · Trained on **Lightning AI** · by [vedaco](https://huggingface.co/vedaco)
+    """)
+
+    with gr.Tabs():
+
+        # ── Tab: Generate ────────────────────────────────────
+        with gr.TabItem("🖼️ Generate"):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    num_images = gr.Slider(1, 16, value=4, step=1, label="Number of images")
+                    temperature = gr.Slider(0.1, 2.0, value=0.9, step=0.05, label="Temperature",
+                                            info="Higher → more diverse")
+                    top_k = gr.Slider(1, 1024, value=100, step=1, label="Top-K")
+                    top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-P")
+                    seed = gr.Number(value=-1, label="Seed (-1 = random)", precision=0)
+                    gen_btn = gr.Button("🎨 Generate", variant="primary", size="lg")
+
+                with gr.Column(scale=2):
+                    gallery = gr.Gallery(label="Generated Images", columns=2, height=512)
+                    gen_info = gr.Textbox(label="Info", interactive=False)
+
+            gen_btn.click(
+                generate_images,
+                inputs=[num_images, temperature, top_k, top_p, seed],
+                outputs=[gallery, gen_info],
+            )
+
+        # ── Tab: Reconstruct ─────────────────────────────────
+        with gr.TabItem("🔄 Reconstruct (VQVAE)"):
+            gr.Markdown("Upload an image to see how the VQVAE tokenizer encodes & decodes it.")
+            with gr.Row():
+                input_img = gr.Image(label="Input", type="numpy")
+                output_img = gr.Image(label="Reconstructed")
+            recon_info = gr.Textbox(label="Info", interactive=False)
+            recon_btn = gr.Button("🔄 Reconstruct", variant="primary")
+            recon_btn.click(reconstruct_image, inputs=[input_img], outputs=[output_img, recon_info])
+
+        # ── Tab: About ────────────────────────────────────────
+        with gr.TabItem("ℹ️ About"):
+            if config:
+                tok = config["model"]["tokenizer"]
+                trans = config["model"]["transformer"]
+                grid = tok["image_size"] // tok["patch_size"]
+                gr.Markdown(f"""
+                ## Architecture
+
+                | Component | Value |
+                |-----------|-------|
+                | Image size | {tok['image_size']}×{tok['image_size']} |
+                | Grid | {grid}×{grid} = {grid*grid} tokens |
+                | Codebook | {tok['num_tokens']} codes × {tok['codebook_dim']}d |
+                | Transformer | {trans['num_layers']} layers, {trans['d_model']}d, {trans['num_heads']} heads |
+                | FFN dim | {trans['d_ff']} |
+                | Position encoding | RoPE |
+                | Activation | SwiGLU |
+                | Normalization | RMSNorm |
+
+                ## How it works
+
+                1. **VQVAE Encoder** compresses a 256×256 image into a {grid}×{grid} grid of discrete tokens
+                2. The grid is flattened into a sequence of **{grid*grid} tokens**
+                3. A **causal transformer** is trained to predict the next token
+                4. At generation time, tokens are sampled one-by-one and decoded with the **VQVAE Decoder**
+                """)
+            else:
+                gr.Markdown("Config not available.")
+
+    gr.Markdown("---\nMade with ❤️ by [vedaco](https://huggingface.co/vedaco) · TensorFlow · Lightning AI")
+
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)