FilesExpand file tree

master

/

latent_preview.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

74 lines (55 loc) · 2.56 KB

master

/

latent_preview.py

Copy path

Top

File metadata and controls

• Code
• Blame

74 lines (55 loc) · 2.56 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

import torch

from PIL import Image

import struct

import numpy as np

from comfy.cli_args import args, LatentPreviewMethod

from comfy.taesd.taesd import TAESD

import folder_paths

MAX_PREVIEW_RESOLUTION = 512

class LatentPreviewer:

def decode_latent_to_preview(self, x0):

pass

def decode_latent_to_preview_image(self, preview_format, x0):

preview_image = self.decode_latent_to_preview(x0)

return ("JPEG", preview_image, MAX_PREVIEW_RESOLUTION)

class TAESDPreviewerImpl(LatentPreviewer):

def __init__(self, taesd):

self.taesd = taesd

def decode_latent_to_preview(self, x0):

x_sample = self.taesd.decoder(x0)[0].detach()

# x_sample = self.taesd.unscale_latents(x_sample).div(4).add(0.5) # returns value in [-2, 2]

x_sample = x_sample.sub(0.5).mul(2)

x_sample = torch.clamp((x_sample + 1.0) / 2.0, min=0.0, max=1.0)

x_sample = 255. * np.moveaxis(x_sample.cpu().numpy(), 0, 2)

x_sample = x_sample.astype(np.uint8)

preview_image = Image.fromarray(x_sample)

return preview_image

class Latent2RGBPreviewer(LatentPreviewer):

def __init__(self, latent_rgb_factors):

self.latent_rgb_factors = torch.tensor(latent_rgb_factors, device="cpu")

def decode_latent_to_preview(self, x0):

latent_image = x0[0].permute(1, 2, 0).cpu() @ self.latent_rgb_factors

latents_ubyte = (((latent_image + 1) / 2)

.clamp(0, 1) # change scale from -1..1 to 0..1

.mul(0xFF) # to 0..255

.byte()).cpu()

return Image.fromarray(latents_ubyte.numpy())

def get_previewer(device, latent_format):

previewer = None

method = args.preview_method

if method != LatentPreviewMethod.NoPreviews:

# TODO previewer methods

taesd_decoder_path = folder_paths.get_full_path("vae_approx", latent_format.taesd_decoder_name)

if method == LatentPreviewMethod.Auto:

method = LatentPreviewMethod.Latent2RGB

if taesd_decoder_path:

method = LatentPreviewMethod.TAESD

if method == LatentPreviewMethod.TAESD:

if taesd_decoder_path:

taesd = TAESD(None, taesd_decoder_path).to(device)

previewer = TAESDPreviewerImpl(taesd)

else:

print("Warning: TAESD previews enabled, but could not find models/vae_approx/{}".format(latent_format.taesd_decoder_name))

if previewer is None:

previewer = Latent2RGBPreviewer(latent_format.latent_rgb_factors)

return previewer