Shopify/background-replacement

1import torch2from diffusers import StableDiffusionXLControlNetPipeline, ControlNetModel, AutoencoderKL, UniPCMultistepScheduler3 4device = None5pipe = None6 7 8def init():9    global device, pipe10 11    device = "cuda" if torch.cuda.is_available() else "cpu"12 13    print("Initializing depth ControlNet...")14 15    depth_controlnet = ControlNetModel.from_pretrained(16        "diffusers/controlnet-depth-sdxl-1.0",17        use_safetensors=True,18        torch_dtype=torch.float1619    ).to(device)20 21    print("Initializing autoencoder...")22 23    vae = AutoencoderKL.from_pretrained(24        "madebyollin/sdxl-vae-fp16-fix",25        torch_dtype=torch.float16,26    ).to(device)27 28    print("Initializing SDXL pipeline...")29 30    pipe = StableDiffusionXLControlNetPipeline.from_pretrained(31        "stabilityai/stable-diffusion-xl-base-1.0",32        controlnet=[depth_controlnet],33        vae=vae,34        variant="fp16",35        use_safetensors=True,36        torch_dtype=torch.float1637        # low_cpu_mem_usage=True38    ).to(device)39 40    pipe.enable_model_cpu_offload()41    # speed up diffusion process with faster scheduler and memory optimization42    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)43    # remove following line if xformers is not installed44    pipe.enable_xformers_memory_efficient_attention()45 46 47def run_pipeline(image, positive_prompt, negative_prompt, seed):48    if seed == -1:49        print("Using random seed")50        generator = None51    else:52        print("Using seed:", seed)53        generator = torch.manual_seed(seed)54 55    images = pipe(56        prompt=positive_prompt,57        negative_prompt=negative_prompt,58        num_inference_steps=30,59        num_images_per_prompt=4,60        controlnet_conditioning_scale=0.65,61        guidance_scale=10.0,62        generator=generator,63        image=image64    ).images65 66    return images67