Upload from local env

app.py

@@ -0,0 +1,245 @@
+from typing import List, Literal
+import gradio as gr
+import torch
+import numpy as np
+import colorsys
+
+from diffusers import VQModel
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.pipelines.wuerstchen.modeling_paella_vq_model import PaellaVQModel
+from abc import abstractmethod
+import torch.backends
+import torch.mps
+from PIL import Image
+
+
+if torch.cuda.is_available():
+    device = torch.device("cuda")
+elif torch.backends.mps.is_available():
+    device = torch.device("mps")
+else:
+    device = torch.device("cpu")
+
+
+# abstract class VQImageRoundtripPipeline:
+class ImageRoundtripPipeline:
+    @abstractmethod
+    def roundtrip_image(self, image, output_type="pil"): ...
+
+
+class VQImageRoundtripPipeline(ImageRoundtripPipeline):
+    vqvae: VQModel
+    vae_scale_factor: int
+    vqvae_processor: VaeImageProcessor
+
+    def __init__(self):
+        self.vqvae = VQModel.from_pretrained("amused/amused-512", subfolder="vqvae")
+        self.vqvae.eval()
+        self.vqvae.to(device)
+        self.vae_scale_factor = 2 ** (len(self.vqvae.config.block_out_channels) - 1)
+        self.vqvae_processor = VaeImageProcessor(
+            vae_scale_factor=self.vae_scale_factor, do_normalize=False
+        )
+        print("VQ-GAN model loaded", self.vqvae)
+
+    def roundtrip_image(self, image, output_type="pil"):
+        image = self.vqvae_processor.preprocess(image)
+        device = self.vqvae.device
+        needs_upcasting = (
+            self.vqvae.dtype == torch.float16 and self.vqvae.config.force_upcast
+        )
+
+        batch_size, im_channels, height, width = image.shape
+
+        if needs_upcasting:
+            self.vqvae.float()
+
+        latents = self.vqvae.encode(
+            image.to(dtype=self.vqvae.dtype, device=device)
+        ).latents
+        latents_batch_size, latent_channels, latents_height, latents_width = (
+            latents.shape
+        )
+        latents = self.vqvae.quantize(latents)[2][2].reshape(
+            batch_size, latents_height, latents_width
+        )
+        output = self.vqvae.decode(
+            latents,
+            force_not_quantize=True,
+            shape=(
+                batch_size,
+                height // self.vae_scale_factor,
+                width // self.vae_scale_factor,
+                self.vqvae.config.latent_channels,
+            ),
+        ).sample.clip(0, 1)
+        output = self.vqvae_processor.postprocess(output, output_type)
+
+        if needs_upcasting:
+            self.vqvae.half()
+
+        return output[0], latents.cpu().numpy(), self.vqvae.config.num_vq_embeddings
+
+
+class PaellaImageRoundtripPipeline(ImageRoundtripPipeline):
+    vqgan: PaellaVQModel
+    vae_scale_factor: int
+    vqvae_processor: VaeImageProcessor
+
+    def __init__(self):
+        self.vqgan = PaellaVQModel.from_pretrained(
+            "warp-ai/wuerstchen", subfolder="vqgan"
+        )
+        self.vqgan.eval()
+        self.vqgan.to(device)
+        self.vae_scale_factor = 4
+        self.vqvae_processor = VaeImageProcessor(
+            vae_scale_factor=self.vae_scale_factor, do_normalize=False
+        )
+        print("Paella VQ-GAN model loaded", self.vqgan)
+
+    def roundtrip_image(self, image, output_type="pil"):
+        image = self.vqvae_processor.preprocess(image)
+        device = self.vqgan.device
+
+        batch_size, im_channels, height, width = image.shape
+
+        latents = self.vqgan.encode(
+            image.to(dtype=self.vqgan.dtype, device=device)
+        ).latents
+        latents_batch_size, latent_channels, latents_height, latents_width = (
+            latents.shape
+        )
+        # latents = latents * self.vqgan.config.scale_factor
+        # Manually quantize so we can inspect
+        latents_q = self.vqgan.vquantizer(latents)[2][2].reshape(
+            batch_size, latents_height, latents_width
+        )
+        print("latents after quantize", (latents_q.shape, latents_q.dtype))
+        images = self.vqgan.decode(latents).sample.clamp(0, 1)
+        output = self.vqvae_processor.postprocess(images, output_type)
+
+        # if needs_upcasting:
+        #     self.vqgan.half()
+
+        return output[0], latents_q.cpu().numpy(), self.vqgan.config.num_vq_embeddings
+
+
+pipeline_paella = PaellaImageRoundtripPipeline()
+pipeline_vq = VQImageRoundtripPipeline()
+
+
+# Function to generate a list of unique colors
+def generate_unique_colors_hsl(n):
+    colors = []
+    for i in range(n):
+        hue = i / (n // 4)  # Distribute hues evenly around the color wheel 4 times
+        lightness = 0.8 - (i / n) * 0.6  # Decrease brightness from 0.8 to 0.2
+        saturation = 1.0
+        rgb = colorsys.hls_to_rgb(hue, lightness, saturation)
+        rgb = tuple(int(255 * x) for x in rgb)
+        colors.append(rgb)
+    return colors
+
+
+# Function to create the image from VQGAN tokens
+def vqgan_tokens_to_image(tokens, codebook_size, downscale_factor):
+    # Generate unique colors for each token in the codebook
+    colors = generate_unique_colors_hsl(codebook_size)
+
+    # Create a lookup table
+    lookup_table = np.array(colors, dtype=np.uint8)
+
+    # Extract the token array (remove the batch dimension)
+    token_array = tokens[0]
+
+    # Map tokens to their RGB colors using the lookup table
+    color_image = lookup_table[token_array]
+
+    # Create a PIL image from the numpy array
+    img = Image.fromarray(color_image, "RGB")
+
+    # Upscale the image using nearest neighbor interpolation
+    img = img.resize(
+        (
+            color_image.shape[1] * downscale_factor,
+            color_image.shape[0] * downscale_factor,
+        ),
+        Image.NEAREST,
+    )
+
+    return img
+
+
+# This is a gradio space that lets you encode an image with various encoder-decoder pairs, eg VQ-GAN, SDXL's VAE, etc and check the image quality
+
+
+# def image_grid_to_string(image_grid):
+#     """Convert a latent vq index "image" grid to a string, input shape is (1, height, width)"""
+#     return "\n".join(
+#         [" ".join([str(int(x)) for x in row]) for row in image_grid.squeeze()]
+#     )
+
+
+def describe_shape(shape):
+    return f"Shape: {shape} num elements: {np.prod(shape)}"
+
+
+# @spaces.GPU
+@torch.no_grad()
+def roundtrip_image(
+    image,
+    model: List[Literal["vqgan", Literal["paella"]]],
+    size: List[Literal["256x256", "512x512", "1024x1024"]],
+    output_type="pil",
+):
+    if size == "256x256":
+        image = image.resize((256, 256))
+    elif size == "512x512":
+        image = image.resize((512, 512))
+    elif size == "1024x1024":
+        image = image.resize((1024, 1024))
+    else:
+        raise ValueError(f"Unknown size {size}")
+
+    if model == "vqgan":
+        image, latents, codebook_size = pipeline_vq.roundtrip_image(image, output_type)
+        return (
+            image,
+            vqgan_tokens_to_image(
+                latents, codebook_size, downscale_factor=pipeline_vq.vae_scale_factor
+            ),
+            describe_shape(latents.shape),
+        )
+    elif model == "paella":
+        image, latents, codebook_size = pipeline_paella.roundtrip_image(
+            image, output_type
+        )
+        return (
+            image,
+            vqgan_tokens_to_image(
+                latents, codebook_size, downscale_factor=pipeline_vq.vae_scale_factor
+            ),
+            describe_shape(latents.shape),
+        )
+    else:
+        raise ValueError(f"Unknown model {model}")
+
+
+demo = gr.Interface(
+    fn=roundtrip_image,
+    inputs=[
+        gr.Image(type="pil"),
+        gr.Dropdown(["vqgan", "paella"], label="Model", value="vqgan"),
+        gr.Dropdown(["256x256", "512x512", "1024x1024"], label="Size", value="512x512"),
+    ],
+    outputs=[
+        gr.Image(label="Reconstructed"),
+        gr.Image(label="Tokens"),
+        gr.Text(label="VQ Shape"),
+    ],
+    title="Image Tokenizer Playground",
+    description="Round-trip an image through an encode-decoder pair to see the quality loss from the VQ-GAN for image generation, etc.",
+)
+
+demo.launch()