Attempt using safetensors for lightweight memory

app.py

@@ -15,10 +15,15 @@ from cldm.model import create_model, load_state_dict
 
 from huggingface_hub import hf_hub_url, cached_download
 
-REPO_ID = "lllyasviel/ControlNet"
-canny_checkpoint = "models/control_sd15_canny.pth"
-scribble_checkpoint = "models/control_sd15_scribble.pth"
-pose_checkpoint = "models/control_sd15_openpose.pth"
+# REPO_ID = "lllyasviel/ControlNet"
+# canny_checkpoint = "models/control_sd15_canny.pth"
+# scribble_checkpoint = "models/control_sd15_scribble.pth"
+# pose_checkpoint = "models/control_sd15_openpose.pth"
+
+REPO_ID = "webui/ControlNet-modules-safetensors"
+canny_checkpoint = " control_canny-fp16.safetensors"
+scribble_checkpoint = "control_scribble-fp16.safetensors"
+pose_checkpoint = "control_openpose-fp16.safetensors"
 
 canny_model = create_model('./models/cldm_v15.yaml').cpu()
 canny_model.load_state_dict(load_state_dict(cached_download(
@@ -30,7 +35,7 @@ ddim_sampler = DDIMSampler(canny_model)
 pose_model = create_model('./models/cldm_v15.yaml').cpu()
 pose_model.load_state_dict(load_state_dict(cached_download(
     hf_hub_url(REPO_ID, pose_checkpoint)
-), location='cuda'))
+), location='cpu'))
 pose_model = pose_model.cuda()
 ddim_sampler_pose = DDIMSampler(pose_model)
 
@@ -41,6 +46,8 @@ scribble_model.load_state_dict(load_state_dict(cached_download(
 scribble_model = canny_model.cuda()
 ddim_sampler_scribble = DDIMSampler(scribble_model)
 
+save_memory = False
+
 def process(input_image, prompt, input_control, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, scale, seed, eta, low_threshold, high_threshold):
     # TODO: Add other control tasks
     if input_control == "Scribble":
@@ -64,14 +71,24 @@ def process_canny(input_image, prompt, a_prompt, n_prompt, num_samples, image_re
 
         seed_everything(seed)
 
+        if save_memory:
+            canny_model.low_vram_shift(is_diffusing=False)
+
         cond = {"c_concat": [control], "c_crossattn": [canny_model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)]}
         un_cond = {"c_concat": [control], "c_crossattn": [canny_model.get_learned_conditioning([n_prompt] * num_samples)]}
         shape = (4, H // 8, W // 8)
 
+        if save_memory:
+            canny_model.low_vram_shift(is_diffusing=False)
+
         samples, intermediates = ddim_sampler.sample(ddim_steps, num_samples,
                                                      shape, cond, verbose=False, eta=eta,
                                                      unconditional_guidance_scale=scale,
                                                      unconditional_conditioning=un_cond)
+
+        if save_memory:
+            canny_model.low_vram_shift(is_diffusing=False)
+            
         x_samples = canny_model.decode_first_stage(samples)
         x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
 
@@ -92,16 +109,24 @@ def process_scribble(input_image, prompt, a_prompt, n_prompt, num_samples, image
 
         seed_everything(seed)
 
+        if save_memory:
+            scribble_model.low_vram_shift(is_diffusing=False)
+
         cond = {"c_concat": [control], "c_crossattn": [scribble_model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)]}
         un_cond = {"c_concat": [control], "c_crossattn": [scribble_model.get_learned_conditioning([n_prompt] * num_samples)]}
         shape = (4, H // 8, W // 8)
 
-
+        if save_memory:
+            scribble_model.low_vram_shift(is_diffusing=False)
+            
         samples, intermediates = ddim_sampler_scribble.sample(ddim_steps, num_samples,
                                                      shape, cond, verbose=False, eta=eta,
                                                      unconditional_guidance_scale=scale,
                                                      unconditional_conditioning=un_cond)
 
+        if save_memory:
+            scribble_model.low_vram_shift(is_diffusing=False)
+                    
         x_samples = scribble_model.decode_first_stage(samples)
         x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)
 
@@ -126,18 +151,25 @@ def process_pose(input_image, prompt, a_prompt, n_prompt, num_samples, image_res
             seed = random.randint(0, 65535)
         seed_everything(seed)
 
+        if save_memory:
+            pose_model.low_vram_shift(is_diffusing=False)
 
+    
         cond = {"c_concat": [control], "c_crossattn": [pose_model.get_learned_conditioning([prompt + ', ' + a_prompt] * num_samples)]}
         un_cond = {"c_concat": [control], "c_crossattn": [pose_model.get_learned_conditioning([n_prompt] * num_samples)]}
         shape = (4, H // 8, W // 8)
 
-
+        if save_memory:
+            pose_model.low_vram_shift(is_diffusing=False)
+            
         samples, intermediates = ddim_sampler_pose.sample(ddim_steps, num_samples,
                                                      shape, cond, verbose=False, eta=eta,
                                                      unconditional_guidance_scale=scale,
                                                      unconditional_conditioning=un_cond)
 
-
+        if save_memory:
+            pose_model.low_vram_shift(is_diffusing=False)
+            
         x_samples = pose_model.decode_first_stage(samples)
         x_samples = (einops.rearrange(x_samples, 'b c h w -> b h w c') * 127.5 + 127.5).cpu().numpy().clip(0, 255).astype(np.uint8)