AiArtLab/sdxs-1b

Instalación y uso básico con Diffusers:

pip install -U diffusers transformers accelerate

import torch
from diffusers import DiffusionPipeline

# usar "mps" para dispositivos Apple
pipe = DiffusionPipeline.from_pretrained(
    "AiArtLab/sdxs-1b",
    dtype=torch.bfloat16,
    device_map="cuda"
)

prompt = "sdxs-1b"
image = pipe(prompt).images[0]

Ejemplo de texto a imagen:

#!pip install -U torch torchvision
#!pip install -U diffusers accelerate transformers

import torch
from diffusers import DiffusionPipeline

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if torch.cuda.is_available() else torch.float32
pipe_id = "AiArtLab/sdxs-1b"

pipe = DiffusionPipeline.from_pretrained(
    pipe_id,
    torch_dtype=dtype,
    trust_remote_code=True
).to(device)

prompt = "1girl, smiling, red eyes, blue hair, white shirt"
refined = pipe.refine_prompts(prompt)[0]
negative_prompt = "bad quality grainy image with low details, incomplete text, despite numerous technical flaws and distorted figures"

image = pipe(
    prompt=refined,
    negative_prompt=negative_prompt,
    width=1024,
    height=1408,
    guidance_scale=5,
    seed=0,
).images[0]

image.show()

Escalado de imagen con el VAE integrado:

upscaled = pipe.image_upscale("media/girl.jpg")
upscaled[0].show()

Ejemplo independiente de uso del VAE para upscale 2x:

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]:
    images = [image] if isinstance(image, (str, Image.Image)) else image
    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))

    unique_shapes = {t.shape for t, _, _ in batch_data}
    step = batch_size if len(unique_shapes) == 1 else 1
    output_images = []

    for i in range(0, len(batch_data), step):
        chunk = batch_data[i : i + step]
        tensors = torch.stack([c[0] for c in chunk])
        latents = vae.encode(tensors).latent_dist.mean
        decoded = vae.decode(latents.to(vae.dtype))[0]
        decoded = (decoded.clamp(-1, 1) + 1) / 2

        for j, tensor in enumerate(decoded):
            w, h = chunk[j][1], chunk[j][2]
            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()