Create akasha/generate.py

akasha/generate.py

@@ -0,0 +1,109 @@
+"""
+Autoregressive generation for AKASHA.
+Generates images token-by-token using the transformer.
+"""
+
+import tensorflow as tf
+import numpy as np
+from tqdm import tqdm
+
+
+class AKASHAGenerator:
+    """Handles autoregressive image generation."""
+
+    def __init__(self, model, config):
+        self.model = model
+        self.config = config
+        tok = config["model"]["tokenizer"]
+        self.image_size = tok["image_size"]
+        self.patch_size = tok["patch_size"]
+        self.num_tokens = tok["num_tokens"]
+        self.grid_size = self.image_size // self.patch_size
+        self.seq_length = self.grid_size * self.grid_size
+        self.bos_token = self.num_tokens
+        self.eos_token = self.num_tokens + 1
+
+    def generate(self, num_images=1, temperature=0.9, top_k=100,
+                 top_p=0.95, seed=None, callback=None):
+        """Generate images autoregressively."""
+        if seed is not None:
+            tf.random.set_seed(seed)
+
+        current_tokens = tf.cast(
+            tf.fill([num_images, 1], self.bos_token), tf.int32
+        )
+
+        generated_tokens = []
+        desc = f"Generating {num_images} image(s) ({self.seq_length} tokens)"
+
+        for step in tqdm(range(self.seq_length), desc=desc):
+            logits = self.model.transformer(current_tokens, training=False)
+            next_token_logits = logits[:, -1, :] / temperature
+
+            # Top-k filtering
+            if top_k > 0:
+                top_k_actual = min(top_k, next_token_logits.shape[-1])
+                values, _ = tf.math.top_k(next_token_logits, k=top_k_actual)
+                min_value = values[:, -1:]
+                next_token_logits = tf.where(
+                    next_token_logits < min_value,
+                    tf.fill(tf.shape(next_token_logits), -1e10),
+                    next_token_logits,
+                )
+
+            next_token = tf.random.categorical(
+                next_token_logits, num_samples=1, dtype=tf.int32
+            )
+            next_token = tf.clip_by_value(next_token, 0, self.num_tokens - 1)
+
+            generated_tokens.append(next_token)
+            current_tokens = tf.concat([current_tokens, next_token], axis=1)
+
+            if callback:
+                callback(step + 1, self.seq_length)
+
+        token_sequence = tf.concat(generated_tokens, axis=1)
+
+        from akasha.tokenizer import ImageTokenizer
+        tokenizer = ImageTokenizer(self.model.vqvae)
+        images = tokenizer.detokenize(token_sequence, grid_size=self.grid_size)
+
+        return images, token_sequence
+
+    def generate_with_kv_cache(self, num_images=1, temperature=0.9,
+                                top_k=100, top_p=0.95):
+        """Generate with KV cache for faster inference."""
+        current_token = tf.cast(
+            tf.fill([num_images, 1], self.bos_token), tf.int32
+        )
+        past_kvs = [None] * self.model.transformer.num_layers
+        generated_tokens = []
+
+        for step in tqdm(range(self.seq_length), desc="Generating (cached)"):
+            logits, past_kvs = self.model.transformer(
+                current_token, training=False, use_cache=True, past_kvs=past_kvs
+            )
+            next_token_logits = logits[:, -1, :] / temperature
+
+            if top_k > 0:
+                top_k_actual = min(top_k, next_token_logits.shape[-1])
+                values, _ = tf.math.top_k(next_token_logits, k=top_k_actual)
+                min_value = values[:, -1:]
+                next_token_logits = tf.where(
+                    next_token_logits < min_value,
+                    tf.fill(tf.shape(next_token_logits), -1e10),
+                    next_token_logits,
+                )
+
+            next_token = tf.random.categorical(
+                next_token_logits, num_samples=1, dtype=tf.int32
+            )
+            next_token = tf.clip_by_value(next_token, 0, self.num_tokens - 1)
+            generated_tokens.append(next_token)
+            current_token = next_token
+
+        token_sequence = tf.concat(generated_tokens, axis=1)
+        from akasha.tokenizer import ImageTokenizer
+        tokenizer = ImageTokenizer(self.model.vqvae)
+        images = tokenizer.detokenize(token_sequence, grid_size=self.grid_size)
+        return images, token_sequence