| import torch |
| import os |
|
|
| from diffusers import ( |
| DDPMScheduler, |
| StableDiffusionXLImg2ImgPipeline, |
| LTXPipeline, |
| AutoencoderKL, |
| ) |
|
|
| from hidiffusion import apply_hidiffusion |
|
|
| from mediapipe.tasks import python |
| from mediapipe.tasks.python import vision |
|
|
| from image_gen_aux import UpscaleWithModel |
|
|
| BASE_MODEL = "stabilityai/sdxl-turbo" |
| VIDEO_MODEL = "Lightricks/LTX-Video" |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
| class ModelHandler: |
| def __init__(self): |
| self.base_pipe = None |
| self.video_pipe = None |
| self.compiled_model = None |
| self.segmenter = None |
| self.upscaler = None |
| self.upscaler4SD = None |
| self.load_models() |
|
|
| def load_base(self): |
| vae = AutoencoderKL.from_pretrained( |
| "madebyollin/sdxl-vae-fp16-fix", |
| torch_dtype=torch.float16, |
| ) |
|
|
| base_pipe = StableDiffusionXLImg2ImgPipeline.from_pretrained( |
| BASE_MODEL, |
| vae=vae, |
| torch_dtype=torch.float16, |
| variant="fp16", |
| use_safetensors=True, |
| ) |
| base_pipe = base_pipe.to(device, silence_dtype_warnings=True) |
| base_pipe.scheduler = DDPMScheduler.from_pretrained( |
| BASE_MODEL, |
| subfolder="scheduler", |
| ) |
| apply_hidiffusion(base_pipe) |
| |
| return base_pipe |
| |
| def load_video_pipe(self): |
| pipe = LTXPipeline.from_pretrained(VIDEO_MODEL, torch_dtype=torch.bfloat16) |
| pipe.to(device) |
| return pipe |
| |
| def load_segmenter(self): |
| segment_model = "checkpoints/selfie_multiclass_256x256.tflite" |
| base_options = python.BaseOptions(model_asset_path=segment_model) |
| options = vision.ImageSegmenterOptions(base_options=base_options,output_category_mask=True) |
| segmenter = vision.ImageSegmenter.create_from_options(options) |
| return segmenter |
| |
| def load_upscaler(self): |
| model_name = os.environ.get("UPSCALE_MODEL", "Phips/4xNomosWebPhoto_RealPLKSR") |
| upscaler = UpscaleWithModel.from_pretrained(model_name).to(device) |
| return upscaler |
| |
| def load_upscaler4SD(self): |
| model_name = os.environ.get("UPSCALE_FOR_SD_MODEL", "Phips/1xDeJPG_realplksr_otf") |
| upscaler = UpscaleWithModel.from_pretrained(model_name).to(device) |
| return upscaler |
|
|
| def load_models(self): |
| base_pipe = self.load_base() |
| segmenter = self.load_segmenter() |
| upscaler = self.load_upscaler() |
| upscaler4SD = self.load_upscaler4SD() |
|
|
| self.base_pipe = base_pipe |
| self.segmenter = segmenter |
| self.upscaler = upscaler |
| self.upscaler4SD = upscaler4SD |
|
|
| MODELS = ModelHandler() |
