Spaces:

Pixelkram
/

biegen

Paused

biegen / engine_sd15.py

Deploy app

1624300 verified 6 days ago

13 kB

"""Stable Diffusion 1.5 engine with the same interface and rules as the FLUX engine.
The loop runs Euler in sigma space, so the update reads like the FLUX one:
x = x + (sigma_next - sigma) * eps. The network is told the time as a timestep in [0, 1000).
"""

from __future__ import annotations

import threading
from collections import OrderedDict
from contextlib import contextmanager

import torch
from PIL import Image

import engine as E
from engine import Clock, Settings

DECODE_PIXELS = E.DECODE_PIXELS


class SD15Engine(E.Engine):
    latent_factor = 8               # VAE scale, no patching
    standard_layers = (12,)         # CLIP's last layer, as in the pipeline

    def __init__(self, pipe, device: str | torch.device = "cuda"):
        from diffusers import EulerDiscreteScheduler

        self.pipe = pipe
        self.device = torch.device(device)
        self.clock = Clock()
        self._emb_cache: OrderedDict = OrderedDict()
        self._lock = threading.Lock()
        self.post_step: list = []       # functions latent -> latent after every sampler step (Tab 4)
        self.encode_tag = ""            # part of the encoding cache key while a tool changes CLIP
        self.run_notes: list = []
        self.dtype = pipe.unet.dtype
        # Euler in sigma space, configured from the checkpoint's scheduler (same betas, same offset).
        self._scheduler_config = EulerDiscreteScheduler.from_config(pipe.scheduler.config).config
        self._Scheduler = EulerDiscreteScheduler
        for m in (pipe.unet, pipe.text_encoder, pipe.vae):
            m.requires_grad_(False)
            m.eval()

    @classmethod
    def from_pretrained(cls, model_id: str = E.MODEL_ID, device: str = "cuda") -> "SD15Engine":
        from diffusers import StableDiffusionPipeline

        # bfloat16 rather than float16: bent activations overflow less.
        dtype = torch.bfloat16 if (device != "cuda" or torch.cuda.is_bf16_supported()) else torch.float16
        pipe = StableDiffusionPipeline.from_pretrained(model_id, dtype=dtype, variant="fp16", safety_checker=None,
                                                       feature_extractor=None, requires_safety_checker=False)
        pipe = pipe.to(device)
        pipe.vae.to(torch.float32)
        # Typed sizes go up to 1536 px, where the fp32 decoder would run out of memory. SD's VAE would
        # tile from 512 px on, so tiles start only above 1024 px and smaller images decode as before.
        pipe.vae.enable_tiling()
        pipe.vae.tile_sample_min_size = 1024
        pipe.vae.tile_latent_min_size = 1024 // 8
        return cls(pipe, device)

    # ---------------------------------------------------------------- text

    def encode(self, prompt: str, layers: tuple = (12,), max_len: int = 77,
               shuffle_seed: int | None = None) -> torch.Tensor:
        """CLIP hidden states (1, 77, 768) after the final layer norm, read at one layer. Cached.
        shuffle_seed permutes the prompt's tokens before the encoder reads them (R4).
        """
        if len(layers) != 1:
            raise ValueError("SD 1.5 reads the text encoder at exactly one layer")
        key = (prompt, tuple(layers), max_len, shuffle_seed, self.encode_tag)
        with self._lock:
            if key in self._emb_cache:
                self._emb_cache.move_to_end(key)
                return self._emb_cache[key]
        te, tok = self.pipe.text_encoder, self.pipe.tokenizer
        ids = tok(prompt, padding="max_length", max_length=tok.model_max_length, truncation=True,
                  return_tensors="pt").input_ids
        pos = self.content_positions(prompt)
        if shuffle_seed is not None and len(pos) > 1:
            perm = torch.randperm(len(pos), generator=torch.Generator("cpu").manual_seed(shuffle_seed))
            idx = torch.tensor(pos)
            ids[0, idx] = ids[0, idx[perm]]
        ids = ids.to(self.device)
        with torch.inference_mode():
            out = te(ids, output_hidden_states=True)
            # Layer 12 gives exactly the pipeline's prompt_embeds (last_hidden_state).
            # transformers 5 keeps the norm on the model itself, older versions on .text_model.
            norm = te.final_layer_norm if hasattr(te, "final_layer_norm") else te.text_model.final_layer_norm
            emb = norm(out.hidden_states[layers[0]]).to(self.dtype)
        with self._lock:
            self._emb_cache[key] = emb
            while len(self._emb_cache) > E.EMB_CACHE_SIZE:
                self._emb_cache.popitem(last=False)
        return emb

    def prompt_embeddings(self, s: Settings, max_len: int) -> tuple[torch.Tensor, torch.Tensor]:
        # The UNet's cross-attention has no positions on the text side, so permuting the encoded
        # vectors would change nothing. R4 shuffles the word order before CLIP reads it instead.
        seed = s.seed + E.SEED_SHUFFLE if s.emb_shuffle else None
        return (self.encode(s.prompt_a, s.te_layers, max_len, seed).float(),
                self.encode(s.prompt_b, s.te_layers, max_len, seed).float())

    def shuffle_cond(self, s: Settings, cond: torch.Tensor, max_len: int) -> torch.Tensor:
        return cond                                        # already shuffled before the encoder

    def content_positions(self, prompt: str, max_len: int = 77) -> list[int]:
        """Positions of the prompt's tokens: after the start token, before the end token and the padding."""
        n = len(self.pipe.tokenizer(prompt, add_special_tokens=False).input_ids)
        return list(range(1, 1 + min(n, self.pipe.tokenizer.model_max_length - 2)))

    # ---------------------------------------------------------------- blocks

    def block(self, block_id: str):
        unet = self.pipe.unet
        if block_id == "mid":
            return unet.mid_block
        kind, k = block_id.split(".")
        return {"down": unet.down_blocks, "up": unet.up_blocks}[kind][int(k)]

    def block_ids(self) -> list[str]:
        unet = self.pipe.unet
        return ([f"down.{k}" for k in range(len(unet.down_blocks))] + ["mid"]
                + [f"up.{k}" for k in range(len(unet.up_blocks))])

    def reads_text(self, block_id: str) -> bool:
        return bool(getattr(self.block(block_id), "has_cross_attention", False))

    def text_bend_note(self, s: Settings) -> str | None:
        if s.bend_stream == "text" and not self.reads_text(s.bend_block):
            return ("Dieser Block liest den Text nicht, er hat keine Cross-Attention. "
                    "Der Eingriff ändert deshalb nichts. Wählt einen anderen Block.")
        return None

    @contextmanager
    def interventions(self, s: Settings, lat_h: int, lat_w: int, txt_len: int):
        """R8 and R9 as in the FLUX engine, plus the tools of Tab 4."""
        import tools_sd15

        with super().interventions(s, lat_h, lat_w, txt_len), tools_sd15.apply(self, s) as notes:
            self.run_notes += notes
            yield

    @contextmanager
    def bending(self, s: Settings, lat_h: int, lat_w: int, txt_len: int):
        """R8. Image stream: the block's output feature map. Text stream: the text the block attends to."""
        if not s.bend_block or s.bend_op == "none" or s.bend_strength == 0:
            yield
            return
        op, st, image = s.bend_op, float(s.bend_strength), s.bend_stream == "image"
        if op not in E.BEND_OPS:
            raise ValueError(f"unknown bending op {op!r}")
        if not image and op in E.SPATIAL_OPS:
            raise ValueError(f"{op} works on the image stream only")
        if not image and not self.reads_text(s.bend_block):
            yield                                          # nothing to bend, render() adds a note
            return
        block = self.block(s.bend_block)
        masks: dict[int, torch.Tensor] = {}
        gen = torch.Generator(self.device).manual_seed(s.seed + E.SEED_BEND_NOISE) if op == "noise" else None
        clock = self.clock

        def mask(dim: int, device, dtype) -> torch.Tensor | None:
            if op != "ablate":
                return None
            if dim not in masks:                           # fixed per call, one per channel count
                off = torch.randperm(dim, generator=torch.Generator("cpu").manual_seed(s.seed + E.SEED_ABLATE))
                m = torch.ones(dim)
                m[off[:round(min(max(st, 0.0), 1.0) * dim)]] = 0.0   # typed values may leave 0..1
                masks[dim] = m
            return masks[dim].to(device, dtype)

        def bend(h: torch.Tensor) -> torch.Tensor:
            if h.dim() == 4:                               # feature map (B, C, H, W) as tokens (B, H*W, C)
                b, c, hh, ww = h.shape
                tokens = h.permute(0, 2, 3, 1).reshape(b, hh * ww, c)
                out = E.bend_tensor(tokens, op, st, (hh, ww), mask(c, h.device, h.dtype), gen)
                return out.reshape(b, hh, ww, c).permute(0, 3, 1, 2).contiguous()
            return E.bend_tensor(h, op, st, None, mask(h.shape[-1], h.device, h.dtype), gen)

        def active() -> bool:
            return s.bend_from <= clock.progress <= s.bend_to

        if image:
            def hook(module, args, output):
                if not active():
                    return None
                if isinstance(output, tuple):              # down blocks: (hidden states, skip connections)
                    return (bend(output[0]),) + tuple(output[1:])
                return bend(output)

            handle = block.register_forward_hook(hook)
        else:
            def pre_hook(module, args, kwargs):
                if not active() or kwargs.get("encoder_hidden_states") is None:
                    return None
                return args, {**kwargs, "encoder_hidden_states": bend(kwargs["encoder_hidden_states"])}

            handle = block.register_forward_pre_hook(pre_hook, with_kwargs=True)
        try:
            yield
        finally:
            handle.remove()

    # ---------------------------------------------------------------- loop

    def denoise(self, s: Settings, cond: torch.Tensor, uncond: torch.Tensor, seeds: list[int],
                guidances: list[float], on_step=None) -> torch.Tensor:
        """Euler in sigma space for one chunk. Returns latents (n, 4, h, w) in float32."""
        dev, N = self.device, s.steps
        n = len(seeds)
        lat_h, lat_w = s.height // 8, s.width // 8
        noise = torch.stack([torch.randn((4, lat_h, lat_w), generator=torch.Generator("cpu").manual_seed(sd))
                             for sd in seeds]).to(dev)        # float32, device independent
        sched = self._Scheduler.from_config(self._scheduler_config)   # fresh: the scheduler is stateful
        sched.set_timesteps(N, device=dev)
        sig = sched.sigmas.float().cpu().tolist()             # N + 1 values ending in 0
        lat = noise * sched.init_noise_sigma
        n_run = max(1, round(N * s.stop_at))

        enc = torch.cat([cond, uncond.expand(n, -1, -1)]).to(self.dtype)
        g = torch.tensor(guidances, device=dev, dtype=torch.float32).view(-1, 1, 1, 1)
        x0_pred = None
        for i, t in enumerate(sched.timesteps[:n_run]):
            self.clock.progress = i / N
            sigma = sig[i]
            t_model = float(t) * s.sigma_mult + s.sigma_offset * 1000   # R6: the model is told a different time
            x = sched.scale_model_input(lat, t)                         # x / sqrt(sigma^2 + 1)
            x = torch.cat([x, x]).to(self.dtype)
            eps = self.pipe.unet(x, torch.full((2 * n,), t_model, device=dev), encoder_hidden_states=enc,
                                 return_dict=False)[0].float()
            e_c, e_u = eps.chunk(2)
            eps = e_u + g * (e_c - e_u)                                 # R3: per-sample guidance
            x0_pred = lat - sigma * eps
            lat = sched.step(eps, t, lat, return_dict=False)[0]         # x + (sigma_next - sigma) * eps
            for fn in self.post_step:                                   # Tab 4, e.g. the mirror
                lat = fn(lat)
            if on_step is not None:
                on_step()
        return x0_pred if (s.show_x0 and n_run < N) else lat

    def decode(self, final: torch.Tensor, lat_h: int, lat_w: int) -> tuple[list[Image.Image], list[bool]]:
        vae = self.pipe.vae
        n = final.shape[0]
        finite = torch.isfinite(final).flatten(1).all(1).tolist()
        lat = torch.nan_to_num(final.float(), 0.0, 0.0, 0.0) / vae.config.scaling_factor
        per_call = max(1, DECODE_PIXELS // (lat_h * lat_w * 64))
        images = []
        for k in range(0, n, per_call):
            img = vae.decode(lat[k:k + per_call].to(vae.dtype), return_dict=False)[0]
            img = torch.nan_to_num(img.float(), 0.0, 0.0, 0.0)
            images += self.pipe.image_processor.postprocess(img, output_type="pil",
                                                           do_denormalize=[True] * img.shape[0])
        return images, finite

	"""Stable Diffusion 1.5 engine with the same interface and rules as the FLUX engine.

	The loop runs Euler in sigma space, so the update reads like the FLUX one:
	x = x + (sigma_next - sigma) * eps. The network is told the time as a timestep in [0, 1000).
	"""

	from __future__ import annotations

	import threading
	from collections import OrderedDict
	from contextlib import contextmanager

	import torch
	from PIL import Image

	import engine as E
	from engine import Clock, Settings

	DECODE_PIXELS = E.DECODE_PIXELS


	class SD15Engine(E.Engine):
	latent_factor = 8 # VAE scale, no patching
	standard_layers = (12,) # CLIP's last layer, as in the pipeline

	def __init__(self, pipe, device: str \| torch.device = "cuda"):
	from diffusers import EulerDiscreteScheduler

	self.pipe = pipe
	self.device = torch.device(device)
	self.clock = Clock()
	self._emb_cache: OrderedDict = OrderedDict()
	self._lock = threading.Lock()
	self.post_step: list = [] # functions latent -> latent after every sampler step (Tab 4)
	self.encode_tag = "" # part of the encoding cache key while a tool changes CLIP
	self.run_notes: list = []
	self.dtype = pipe.unet.dtype
	# Euler in sigma space, configured from the checkpoint's scheduler (same betas, same offset).
	self._scheduler_config = EulerDiscreteScheduler.from_config(pipe.scheduler.config).config
	self._Scheduler = EulerDiscreteScheduler
	for m in (pipe.unet, pipe.text_encoder, pipe.vae):
	m.requires_grad_(False)
	m.eval()

	@classmethod
	def from_pretrained(cls, model_id: str = E.MODEL_ID, device: str = "cuda") -> "SD15Engine":
	from diffusers import StableDiffusionPipeline

	# bfloat16 rather than float16: bent activations overflow less.
	dtype = torch.bfloat16 if (device != "cuda" or torch.cuda.is_bf16_supported()) else torch.float16
	pipe = StableDiffusionPipeline.from_pretrained(model_id, dtype=dtype, variant="fp16", safety_checker=None,
	feature_extractor=None, requires_safety_checker=False)
	pipe = pipe.to(device)
	pipe.vae.to(torch.float32)
	# Typed sizes go up to 1536 px, where the fp32 decoder would run out of memory. SD's VAE would
	# tile from 512 px on, so tiles start only above 1024 px and smaller images decode as before.
	pipe.vae.enable_tiling()
	pipe.vae.tile_sample_min_size = 1024
	pipe.vae.tile_latent_min_size = 1024 // 8
	return cls(pipe, device)

	# ---------------------------------------------------------------- text

	def encode(self, prompt: str, layers: tuple = (12,), max_len: int = 77,
	shuffle_seed: int \| None = None) -> torch.Tensor:
	"""CLIP hidden states (1, 77, 768) after the final layer norm, read at one layer. Cached.

	shuffle_seed permutes the prompt's tokens before the encoder reads them (R4).
	"""
	if len(layers) != 1:
	raise ValueError("SD 1.5 reads the text encoder at exactly one layer")
	key = (prompt, tuple(layers), max_len, shuffle_seed, self.encode_tag)
	with self._lock:
	if key in self._emb_cache:
	self._emb_cache.move_to_end(key)
	return self._emb_cache[key]
	te, tok = self.pipe.text_encoder, self.pipe.tokenizer
	ids = tok(prompt, padding="max_length", max_length=tok.model_max_length, truncation=True,
	return_tensors="pt").input_ids
	pos = self.content_positions(prompt)
	if shuffle_seed is not None and len(pos) > 1:
	perm = torch.randperm(len(pos), generator=torch.Generator("cpu").manual_seed(shuffle_seed))
	idx = torch.tensor(pos)
	ids[0, idx] = ids[0, idx[perm]]
	ids = ids.to(self.device)
	with torch.inference_mode():
	out = te(ids, output_hidden_states=True)
	# Layer 12 gives exactly the pipeline's prompt_embeds (last_hidden_state).
	# transformers 5 keeps the norm on the model itself, older versions on .text_model.
	norm = te.final_layer_norm if hasattr(te, "final_layer_norm") else te.text_model.final_layer_norm
	emb = norm(out.hidden_states[layers[0]]).to(self.dtype)
	with self._lock:
	self._emb_cache[key] = emb
	while len(self._emb_cache) > E.EMB_CACHE_SIZE:
	self._emb_cache.popitem(last=False)
	return emb

	def prompt_embeddings(self, s: Settings, max_len: int) -> tuple[torch.Tensor, torch.Tensor]:
	# The UNet's cross-attention has no positions on the text side, so permuting the encoded
	# vectors would change nothing. R4 shuffles the word order before CLIP reads it instead.
	seed = s.seed + E.SEED_SHUFFLE if s.emb_shuffle else None
	return (self.encode(s.prompt_a, s.te_layers, max_len, seed).float(),
	self.encode(s.prompt_b, s.te_layers, max_len, seed).float())

	def shuffle_cond(self, s: Settings, cond: torch.Tensor, max_len: int) -> torch.Tensor:
	return cond # already shuffled before the encoder

	def content_positions(self, prompt: str, max_len: int = 77) -> list[int]:
	"""Positions of the prompt's tokens: after the start token, before the end token and the padding."""
	n = len(self.pipe.tokenizer(prompt, add_special_tokens=False).input_ids)
	return list(range(1, 1 + min(n, self.pipe.tokenizer.model_max_length - 2)))

	# ---------------------------------------------------------------- blocks

	def block(self, block_id: str):
	unet = self.pipe.unet
	if block_id == "mid":
	return unet.mid_block
	kind, k = block_id.split(".")
	return {"down": unet.down_blocks, "up": unet.up_blocks}[kind][int(k)]

	def block_ids(self) -> list[str]:
	unet = self.pipe.unet
	return ([f"down.{k}" for k in range(len(unet.down_blocks))] + ["mid"]
	+ [f"up.{k}" for k in range(len(unet.up_blocks))])

	def reads_text(self, block_id: str) -> bool:
	return bool(getattr(self.block(block_id), "has_cross_attention", False))

	def text_bend_note(self, s: Settings) -> str \| None:
	if s.bend_stream == "text" and not self.reads_text(s.bend_block):
	return ("Dieser Block liest den Text nicht, er hat keine Cross-Attention. "
	"Der Eingriff ändert deshalb nichts. Wählt einen anderen Block.")
	return None

	@contextmanager
	def interventions(self, s: Settings, lat_h: int, lat_w: int, txt_len: int):
	"""R8 and R9 as in the FLUX engine, plus the tools of Tab 4."""
	import tools_sd15

	with super().interventions(s, lat_h, lat_w, txt_len), tools_sd15.apply(self, s) as notes:
	self.run_notes += notes
	yield

	@contextmanager
	def bending(self, s: Settings, lat_h: int, lat_w: int, txt_len: int):
	"""R8. Image stream: the block's output feature map. Text stream: the text the block attends to."""
	if not s.bend_block or s.bend_op == "none" or s.bend_strength == 0:
	yield
	return
	op, st, image = s.bend_op, float(s.bend_strength), s.bend_stream == "image"
	if op not in E.BEND_OPS:
	raise ValueError(f"unknown bending op {op!r}")
	if not image and op in E.SPATIAL_OPS:
	raise ValueError(f"{op} works on the image stream only")
	if not image and not self.reads_text(s.bend_block):
	yield # nothing to bend, render() adds a note
	return
	block = self.block(s.bend_block)
	masks: dict[int, torch.Tensor] = {}
	gen = torch.Generator(self.device).manual_seed(s.seed + E.SEED_BEND_NOISE) if op == "noise" else None
	clock = self.clock

	def mask(dim: int, device, dtype) -> torch.Tensor \| None:
	if op != "ablate":
	return None
	if dim not in masks: # fixed per call, one per channel count
	off = torch.randperm(dim, generator=torch.Generator("cpu").manual_seed(s.seed + E.SEED_ABLATE))
	m = torch.ones(dim)
	m[off[:round(min(max(st, 0.0), 1.0) * dim)]] = 0.0 # typed values may leave 0..1
	masks[dim] = m
	return masks[dim].to(device, dtype)

	def bend(h: torch.Tensor) -> torch.Tensor:
	if h.dim() == 4: # feature map (B, C, H, W) as tokens (B, H*W, C)
	b, c, hh, ww = h.shape
	tokens = h.permute(0, 2, 3, 1).reshape(b, hh * ww, c)
	out = E.bend_tensor(tokens, op, st, (hh, ww), mask(c, h.device, h.dtype), gen)
	return out.reshape(b, hh, ww, c).permute(0, 3, 1, 2).contiguous()
	return E.bend_tensor(h, op, st, None, mask(h.shape[-1], h.device, h.dtype), gen)

	def active() -> bool:
	return s.bend_from <= clock.progress <= s.bend_to

	if image:
	def hook(module, args, output):
	if not active():
	return None
	if isinstance(output, tuple): # down blocks: (hidden states, skip connections)
	return (bend(output[0]),) + tuple(output[1:])
	return bend(output)

	handle = block.register_forward_hook(hook)
	else:
	def pre_hook(module, args, kwargs):
	if not active() or kwargs.get("encoder_hidden_states") is None:
	return None
	return args, {**kwargs, "encoder_hidden_states": bend(kwargs["encoder_hidden_states"])}

	handle = block.register_forward_pre_hook(pre_hook, with_kwargs=True)
	try:
	yield
	finally:
	handle.remove()

	# ---------------------------------------------------------------- loop

	def denoise(self, s: Settings, cond: torch.Tensor, uncond: torch.Tensor, seeds: list[int],
	guidances: list[float], on_step=None) -> torch.Tensor:
	"""Euler in sigma space for one chunk. Returns latents (n, 4, h, w) in float32."""
	dev, N = self.device, s.steps
	n = len(seeds)
	lat_h, lat_w = s.height // 8, s.width // 8
	noise = torch.stack([torch.randn((4, lat_h, lat_w), generator=torch.Generator("cpu").manual_seed(sd))
	for sd in seeds]).to(dev) # float32, device independent
	sched = self._Scheduler.from_config(self._scheduler_config) # fresh: the scheduler is stateful
	sched.set_timesteps(N, device=dev)
	sig = sched.sigmas.float().cpu().tolist() # N + 1 values ending in 0
	lat = noise * sched.init_noise_sigma
	n_run = max(1, round(N * s.stop_at))

	enc = torch.cat([cond, uncond.expand(n, -1, -1)]).to(self.dtype)
	g = torch.tensor(guidances, device=dev, dtype=torch.float32).view(-1, 1, 1, 1)
	x0_pred = None
	for i, t in enumerate(sched.timesteps[:n_run]):
	self.clock.progress = i / N
	sigma = sig[i]
	t_model = float(t) * s.sigma_mult + s.sigma_offset * 1000 # R6: the model is told a different time
	x = sched.scale_model_input(lat, t) # x / sqrt(sigma^2 + 1)
	x = torch.cat([x, x]).to(self.dtype)
	eps = self.pipe.unet(x, torch.full((2 * n,), t_model, device=dev), encoder_hidden_states=enc,
	return_dict=False)[0].float()
	e_c, e_u = eps.chunk(2)
	eps = e_u + g * (e_c - e_u) # R3: per-sample guidance
	x0_pred = lat - sigma * eps
	lat = sched.step(eps, t, lat, return_dict=False)[0] # x + (sigma_next - sigma) * eps
	for fn in self.post_step: # Tab 4, e.g. the mirror
	lat = fn(lat)
	if on_step is not None:
	on_step()
	return x0_pred if (s.show_x0 and n_run < N) else lat

	def decode(self, final: torch.Tensor, lat_h: int, lat_w: int) -> tuple[list[Image.Image], list[bool]]:
	vae = self.pipe.vae
	n = final.shape[0]
	finite = torch.isfinite(final).flatten(1).all(1).tolist()
	lat = torch.nan_to_num(final.float(), 0.0, 0.0, 0.0) / vae.config.scaling_factor
	per_call = max(1, DECODE_PIXELS // (lat_h * lat_w * 64))
	images = []
	for k in range(0, n, per_call):
	img = vae.decode(lat[k:k + per_call].to(vae.dtype), return_dict=False)[0]
	img = torch.nan_to_num(img.float(), 0.0, 0.0, 0.0)
	images += self.pipe.image_processor.postprocess(img, output_type="pil",
	do_denormalize=[True] * img.shape[0])
	return images, finite