import logging import torch import numpy as np from PIL import Image from diffusers.utils import check_min_version from pipeline import LotusGPipeline, LotusDPipeline from utils.image_utils import colorize_depth_map from contextlib import nullcontext check_min_version('0.28.0.dev0') def load_models(task_name, device): if task_name == 'depth': model_g = 'jingheya/lotus-depth-g-v1-0' model_d = 'jingheya/lotus-depth-d-v1-1' else: model_g = 'jingheya/lotus-normal-g-v1-0' model_d = 'jingheya/lotus-normal-d-v1-0' dtype = torch.float16 pipe_g = LotusGPipeline.from_pretrained( model_g, torch_dtype=dtype, ) pipe_d = LotusDPipeline.from_pretrained( model_d, torch_dtype=dtype, ) pipe_g.to(device) pipe_d.to(device) pipe_g.set_progress_bar_config(disable=True) pipe_d.set_progress_bar_config(disable=True) logging.info(f"Successfully loaded pipelines from {model_g} and {model_d}.") return pipe_g, pipe_d def infer_pipe(pipe, images_batch, task_name, seed, device): if seed is None: generator = None else: generator = torch.Generator(device=device).manual_seed(seed) if torch.backends.mps.is_available(): autocast_ctx = nullcontext() else: autocast_ctx = torch.autocast(pipe.device.type) with autocast_ctx: # Convert list of images to tensor images = [np.array(img.convert('RGB')).astype(np.float16) for img in images_batch] test_images = torch.stack([torch.tensor(img).permute(2, 0, 1) for img in images]) test_images = test_images / 127.5 - 1.0 test_images = test_images.to(device) task_emb = torch.tensor([1, 0]).float().unsqueeze(0).to(device) task_emb = torch.cat([torch.sin(task_emb), torch.cos(task_emb)], dim=-1) task_emb = task_emb.repeat(len(test_images), 1) # Run inference preds = pipe( rgb_in=test_images, prompt='', num_inference_steps=1, generator=generator, output_type='np', timesteps=[999], task_emb=task_emb, ).images # Post-process predictions outputs = [] if task_name == 'depth': for p in preds: output_npy = p.mean(axis=-1) output_color = colorize_depth_map(output_npy) outputs.append(output_color) else: for p in preds: output_npy = p output_color = Image.fromarray((output_npy * 255).astype(np.uint8)) outputs.append(output_color) return outputs def lotus(images_batch, task_name, seed, device, pipe_g, pipe_d): output_d = infer_pipe(pipe_d, images_batch, task_name, seed, device) return output_d # Only returning depth outputs for this application