|
import torch |
|
from typing import List |
|
from dataclasses import dataclass |
|
from gradio_app.utils import rgba_to_rgb |
|
from custum_3d_diffusion.trainings.config_classes import ExprimentConfig, TrainerSubConfig |
|
from custum_3d_diffusion import modules |
|
from custum_3d_diffusion.custum_modules.unifield_processor import AttnConfig, ConfigurableUNet2DConditionModel |
|
from custum_3d_diffusion.trainings.base import BasicTrainer |
|
from custum_3d_diffusion.trainings.utils import load_config |
|
|
|
|
|
@dataclass |
|
class FakeAccelerator: |
|
device: torch.device = torch.device("cuda") |
|
|
|
|
|
def init_trainers(cfg_path: str, weight_dtype: torch.dtype, extras: dict): |
|
accelerator = FakeAccelerator() |
|
cfg: ExprimentConfig = load_config(ExprimentConfig, cfg_path, extras) |
|
init_config: AttnConfig = load_config(AttnConfig, cfg.init_config) |
|
configurable_unet = ConfigurableUNet2DConditionModel(init_config, weight_dtype) |
|
configurable_unet.enable_xformers_memory_efficient_attention() |
|
trainer_cfgs: List[TrainerSubConfig] = [load_config(TrainerSubConfig, trainer) for trainer in cfg.trainers] |
|
trainers: List[BasicTrainer] = [modules.find(trainer.trainer_type)(accelerator, None, configurable_unet, trainer.trainer, weight_dtype, i) for i, trainer in enumerate(trainer_cfgs)] |
|
return trainers, configurable_unet |
|
|
|
from gradio_app.utils import make_image_grid, split_image |
|
def process_image(function, img, guidance_scale=2., merged_image=False, remove_bg=True): |
|
from rembg import remove |
|
if remove_bg: |
|
img = remove(img) |
|
img = rgba_to_rgb(img) |
|
if merged_image: |
|
img = split_image(img, rows=2) |
|
images = function( |
|
image=img, |
|
guidance_scale=guidance_scale, |
|
) |
|
if len(images) > 1: |
|
return make_image_grid(images, rows=2) |
|
else: |
|
return images[0] |
|
|
|
|
|
def process_text(trainer, pipeline, img, guidance_scale=2.): |
|
pipeline.cfg.validation_prompts = [img] |
|
titles, images = trainer.batched_validation_forward(pipeline, guidance_scale=[guidance_scale]) |
|
return images[0] |
|
|
|
|
|
def load_pipeline(config_path, ckpt_path, pipeline_filter=lambda x: True, weight_dtype = torch.bfloat16): |
|
training_config = config_path |
|
load_from_checkpoint = ckpt_path |
|
extras = [] |
|
device = "cuda" |
|
trainers, configurable_unet = init_trainers(training_config, weight_dtype, extras) |
|
shared_modules = dict() |
|
for trainer in trainers: |
|
shared_modules = trainer.init_shared_modules(shared_modules) |
|
|
|
if load_from_checkpoint is not None: |
|
state_dict = torch.load(load_from_checkpoint, map_location="cpu") |
|
configurable_unet.unet.load_state_dict(state_dict, strict=False) |
|
|
|
configurable_unet.unet.to(device, dtype=weight_dtype) |
|
|
|
pipeline = None |
|
trainer_out = None |
|
for trainer in trainers: |
|
if pipeline_filter(trainer.cfg.trainer_name): |
|
pipeline = trainer.construct_pipeline(shared_modules, configurable_unet.unet) |
|
pipeline.set_progress_bar_config(disable=False) |
|
trainer_out = trainer |
|
pipeline = pipeline.to(device, dtype=weight_dtype) |
|
return trainer_out, pipeline |
