readme

README.md

@@ -104,6 +104,83 @@ Development Note: We have not provided direct comparisons with other popular ups
 ### sdxs / swinir2
 <img src="media/0000.png" height="512"/>
 
+### upscale code example
+
+```
+import torch
+import numpy as np
+
+from diffusers import AsymmetricAutoencoderKL
+from typing import List, Union
+from PIL import Image
+
+vae = AsymmetricAutoencoderKL.from_pretrained("AiArtLab/sdxs-1b",subfolder="vae").cuda().half()
+device = "cuda"
+
+@torch.no_grad()
+def image_upscale(
+    image: Union[str, Image.Image, List[Union[str, Image.Image]]], 
+    batch_size: int = 1
+) -> List[Image.Image]:
+    """
+    Upscales images using asymmetric VAE (x2).
+    Uses smart batching: processes in parallel if sizes match, else falls back to sequential.
+    """
+    images = [image] if isinstance(image, (str, Image.Image)) else image
+
+    # 1. Preprocess: Load, Handle Alpha, Pad to %8, Normalize
+    batch_data = []
+    for img in images:
+        if isinstance(img, str): img = Image.open(img)
+        if img.mode == "RGBA":
+            img = Image.alpha_composite(Image.new("RGBA", img.size, (255, 255, 255)), img)
+        img = img.convert("RGB")
+
+        w, h = img.size
+        pw, ph = (8 - w % 8) % 8, (8 - h % 8) % 8
+        if pw or ph:
+            padded = Image.new("RGB", (w + pw, h + ph), (255, 255, 255))
+            padded.paste(img)
+            img = padded
+
+        t = torch.from_numpy(np.array(img).astype(np.float32) / 127.5 - 1.0).permute(2, 0, 1)
+        batch_data.append((t.to(device, torch.float16), w, h))
+
+    # 2. Determine Execution Strategy
+    # If all shapes are identical, use batch_size. Else fallback to 1.
+    unique_shapes = {t.shape for t, _, _ in batch_data}
+    step = batch_size if len(unique_shapes) == 1 else 1
+
+    output_images = []
+
+    # 3. Process Batches
+    for i in range(0, len(batch_data), step):
+        chunk = batch_data[i : i + step]
+
+        # Stack tensors [B, C, H, W]
+        tensors = torch.stack([c[0] for c in chunk])
+
+        # Encode -> Decode (using mean for deterministic upscale)
+        latents = vae.encode(tensors).latent_dist.mean
+        decoded = vae.decode(latents.to(vae.dtype))[0]
+
+        # 4. Post-process: Denormalize and Crop
+        decoded = (decoded.clamp(-1, 1) + 1) / 2
+        for j, tensor in enumerate(decoded):
+            w, h = chunk[j][1], chunk[j][2] # Original sizes
+
+            # Crop to exact 2x
+            arr = tensor.cpu().permute(1, 2, 0).float().numpy()
+            arr = arr[:h * 2, :w * 2]
+
+            output_images.append(Image.fromarray((arr * 255).astype("uint8")))
+
+    return output_images
+    
+up = image_upscale("1.jpg")
+up[0].show()
+```
+
 ### Prompt refine
 ```
 refined = pipe.refine_prompts("girl")