Duplicate from gustproof/sd-vae-roundtrip

.gitattributes

@@ -0,0 +1,34 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,13 @@
+---
+title: Sd Vae Roundtrip
+emoji: 📊
+colorFrom: red
+colorTo: yellow
+sdk: gradio
+sdk_version: 3.28.0
+app_file: app.py
+pinned: false
+duplicated_from: gustproof/sd-vae-roundtrip
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,131 @@
+import gradio as gr
+
+from transformers import CLIPTextModel, CLIPTokenizer
+from diffusers import AutoencoderKL, UNet2DConditionModel, PNDMScheduler
+from PIL import Image
+import PIL
+import torch
+import numpy as np
+
+model_path = "Linaqruf/anything-v3.0"
+vae = AutoencoderKL.from_pretrained(model_path, subfolder="vae")
+print(f"vae loaded from {model_path}")
+
+
+def snap(w, h, d=64, area=640 * 640):
+    s = min(1.0, (area / w / h) ** 0.5)
+    err = lambda a, b: 1 - min(a, b) / max(a, b)
+    sw, sh = map(lambda x: int((x * s) // d * d), (w, h))
+    return min(
+        (
+            (ww, hh)
+            for ww, hh in [(sw, sh), (sw, sh + d), (sw + d, sh), (sw + d, sh + d)]
+            if ww * hh <= area
+        ),
+        key=lambda wh: err(w / h, wh[0] / wh[1]),
+    )
+
+
+def center_crop_image(image, hx, wx):
+    # Get the original image dimensions (HxW)
+    original_width, original_height = image.size
+
+    # Calculate the coordinates for center cropping
+    if original_width / original_height > wx / hx:
+        ww = original_height * wx / hx
+        left, right, top, bottom = (
+            (original_width - ww) / 2,
+            (original_width + ww) / 2,
+            0,
+            original_height,
+        )
+    else:
+        hh = original_width * hx / wx
+        left, right, top, bottom = (
+            0,
+            original_width,
+            (original_height - hh) / 2,
+            (original_height + hh) / 2,
+        )
+    # Crop the image
+    cropped_image = image.crop((left, top, right, bottom))
+
+    # Resize the cropped image to the target size (hxw)
+    cropped_image = cropped_image.resize((wx, hx), Image.Resampling.LANCZOS)
+
+    return cropped_image
+
+
+def preprocess(image):
+    if isinstance(image, torch.Tensor):
+        return image
+    elif isinstance(image, PIL.Image.Image):
+        image = [image]
+
+    if isinstance(image[0], PIL.Image.Image):
+        image = [np.array(i)[None, :] for i in image]
+        image = np.concatenate(image, axis=0)
+        image = np.array(image).astype(np.float32) / 255.0
+        image = image.transpose(0, 3, 1, 2)
+        image = 2.0 * image - 1.0
+        image = torch.from_numpy(image)
+    elif isinstance(image[0], torch.Tensor):
+        image = torch.cat(image, dim=0)
+    return image
+
+
+def numpy_to_pil(images):
+    """
+    Convert a numpy image or a batch of images to a PIL image.
+    """
+    if images.ndim == 3:
+        images = images[None, ...]
+    images = (images * 255).round().astype("uint8")
+    if images.shape[-1] == 1:
+        # special case for grayscale (single channel) images
+        pil_images = [Image.fromarray(image.squeeze(), mode="L") for image in images]
+    else:
+        pil_images = [Image.fromarray(image) for image in images]
+
+    return pil_images
+
+
+def postprocess_image(sample: torch.FloatTensor, output_type: str = "pil"):
+    if output_type not in ["pt", "np", "pil"]:
+        raise ValueError(
+            f"output_type={output_type} is not supported. Make sure to choose one of ['pt', 'np', or 'pil']"
+        )
+
+    # Equivalent to diffusers.VaeImageProcessor.denormalize
+    sample = (sample / 2 + 0.5).clamp(0, 1)
+    if output_type == "pt":
+        return sample
+
+    # Equivalent to diffusers.VaeImageProcessor.pt_to_numpy
+    sample = sample.cpu().permute(0, 2, 3, 1).numpy()
+    if output_type == "np":
+        return sample
+
+    # Output_type must be 'pil'
+    sample = numpy_to_pil(sample)
+    return sample
+
+
+def vae_roundtrip(image, max_resolution: int):
+    w, h = image.size
+    ww, hh = snap(w, h, area=max_resolution**2)
+    cropped = center_crop_image(image, hh, ww)
+    image = preprocess(cropped)
+    with torch.no_grad():
+        dist = vae.encode(image)[0]
+        res = vae.decode(dist.mean, return_dict=False)[0]
+    return cropped, postprocess_image(res)[0]
+
+
+iface = gr.Interface(
+    fn=vae_roundtrip,
+    inputs=[gr.Image(type="pil"), gr.Slider(384, 1024, step=64, value=640)],
+    outputs=[gr.Image(label="center cropped"), gr.Image(label="after roundtrip")],
+    allow_flagging="never",
+)
+iface.launch()

requirements.txt

@@ -0,0 +1,5 @@
+transformers
+diffusers
+accelerate
+torch
+numpy