inference.py

# %%
import typing as tp

from diffusers import (
    AutoencoderKL,
    UNet2DConditionModel,
    DiffusionPipeline,
    UniPCMultistepScheduler,
)

from src.model.customSD import CustomNetwork

import torch
import numpy as np  # 반환 유형을 위해 numpy 임포트
from tqdm import tqdm


class Tabular2ImagePipeline(DiffusionPipeline):
    unet: UNet2DConditionModel
    vae: AutoencoderKL
    scheduler: UniPCMultistepScheduler
    fcn: CustomNetwork
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    def __init__(
        self,
        unet: UNet2DConditionModel,
        vae: AutoencoderKL,
        scheduler: UniPCMultistepScheduler,
        fcn: CustomNetwork,
    ) -> None:
        super().__init__()
        self.unet = unet.eval().to(self.device)
        self.vae = vae.eval().to(self.device)
        self.scheduler = scheduler
        self.fcn = fcn.eval().to(self.device)

    @torch.no_grad()
    def __call__(
        self,
        tabular_input: str,
        generator: tp.Optional[torch.Generator] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 7.5,
    ) -> np.ndarray:

        # input transform
        tabular_input = tabular_input.split()

        try:
            data_pca = torch.tensor(
                [float(item) for item in tabular_input], dtype=torch.float32
            ).to(self.device)
        except ValueError:
            raise ValueError("input must be str of numbers separated by spaces.")

        latent = self.fcn(data_pca)  # cond, uncond

        noisy_latents = torch.randn(
            (1, self.unet.in_channels, 64, 64),
            generator=generator,
            device=data_pca.device,
        )

        noisy_latents = noisy_latents * self.scheduler.init_noise_sigma

        self.scheduler.set_timesteps(num_inference_steps)

        pbar = tqdm(self.scheduler.timesteps, desc="generating image")

        for t in self.scheduler.timesteps:

            latent_model_input = noisy_latents
            # latent_model_input = torch.cat([noisy_latents] * 2, dim = 0)

            model_pred = self.unet(
                latent_model_input,
                t,
                encoder_hidden_states=latent,
            ).sample

            # noise_uncond, noise_cond = model_pred.chunk(2)
            # model_pred = noise_uncond + guidance_scale * ( noise_cond - noise_uncond )

            noisy_latents = self.scheduler.step(
                model_pred, t, noisy_latents
            ).prev_sample

            pbar.update(1)

        # image = self.vae.decode(noisy_latents / 0.18215).sample
        image = self.vae.decode(
            1 / self.vae.config.scaling_factor * noisy_latents, return_dict=False
        )[0]
        image = (image / 2 + 0.5).clamp(0, 1)
        image = image.cpu().permute(0, 2, 3, 1).float().numpy()

        return image


# %%
import os
from omegaconf import OmegaConf
from PIL import Image
from src.config.config import ConfigExperiment, set_seed

from src.model.customSD import CustomNetwork
from diffusers import (
    DDPMScheduler,
    AutoencoderKL,
    UNet2DConditionModel,
    UniPCMultistepScheduler,
)

import torch

REVISION = None
VARIANT = None

path_infer = "./exp_results/20250916_095309_exp"
path_cfg_infer = os.path.join(path_infer, "config.yaml")

cfg_infer: ConfigExperiment = OmegaConf.load(path_cfg_infer)

vae: AutoencoderKL = AutoencoderKL.from_pretrained(
    cfg_infer.model_custom.path_sd,
    subfolder="vae",
    revision=REVISION,
    variant=VARIANT,
)

noise_scheduler: UniPCMultistepScheduler = UniPCMultistepScheduler.from_pretrained(
    cfg_infer.model_custom.path_sd,
    subfolder="scheduler",
)

path_fcn_tarined = os.path.join(path_infer, "best/best.pt")
customnetwork = CustomNetwork(cfg_infer.model_custom)
customnetwork.load_state_dict(torch.load(path_fcn_tarined)["state_dict_model"])

path_unet_trained = os.path.join(path_infer, "best")
unet: UNet2DConditionModel = UNet2DConditionModel.from_pretrained(
    path_unet_trained,
    subfolder="unet",
    revision=REVISION,
    variant=VARIANT,
)

pipeline_infer = Tabular2ImagePipeline(
    vae=vae,
    scheduler=noise_scheduler,
    fcn=customnetwork,
    unet=unet,
)
# %%
num_seed = 42
set_seed(num_seed)

generator = torch.Generator(device=pipeline_infer.device)

values_input_0 = "0 1 0 0 1"
values_input_1 = "0 0 0 0 1"
values_input_2 = "0 0 0 1 0"
values_input_3 = "0 0 1 0 0"
values_input_4 = "0 1 0 0 0"
values_input_5 = "1 0 0 0 0"

list_values = [values_input_0, values_input_1, values_input_2, values_input_3, values_input_4, values_input_5]

for i, values_input in enumerate(list_values):
	result = pipeline_infer(
		tabular_input=values_input,
		generator=generator,
	)

	result: Image = Image.fromarray((result[0] * 255).astype(np.uint8))
	result.save(f"./img_gen/image_gen_{i}_seed_{num_seed}.png")
# %%