| 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 |
