Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-20 01:31:13 +02:00
Code Issues Packages Projects Releases Wiki Activity

refactor all control processors to support unload and offload

Browse Source
This commit is contained in:
Vladimir Mandic
2024-01-25 11:45:47 -05:00
parent c19f95141a
commit 7aa6876759
27 changed files with 154 additions and 350 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/control/proc/dpt.py
+20 -13
View File
@@ -3,24 +3,29 @@ import numpy as np
import torch
from transformers import AutoImageProcessor, DPTForDepthEstimation
from modules import devices
from modules.shared import opts
image_processor: AutoImageProcessor = None
dpt_model: DPTForDepthEstimation = None
class DPTDetector:
def __call__(self, input_image=None):
global image_processor, dpt_model # pylint: disable=global-statement
from modules.control.processors import cache_dir
if image_processor is None:
image_processor = AutoImageProcessor.from_pretrained("Intel/dpt-large", cache_dir=cache_dir)
if dpt_model is None:
dpt_model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large", cache_dir=cache_dir)
def __init__(self, model=None, processor=None):
self.model = model
self.processor = processor
def __call__(self, input_image=None):
from modules.control.processors import cache_dir
if self.processor is None:
self.processor = AutoImageProcessor.from_pretrained("Intel/dpt-large", cache_dir=cache_dir)
if self.model is None:
self.model = DPTForDepthEstimation.from_pretrained("Intel/dpt-large", cache_dir=cache_dir)
self.model.to(devices.device)
with devices.inference_context():
inputs = image_processor(images=input_image, return_tensors="pt")
outputs = dpt_model(**inputs)
inputs = self.processor(images=input_image, return_tensors="pt")
inputs.to(devices.device)
outputs = self.model(**inputs)
predicted_depth = outputs.predicted_depth
prediction = torch.nn.functional.interpolate(
predicted_depth.unsqueeze(1),
@@ -30,6 +35,8 @@ class DPTDetector:
)
output = prediction.squeeze().cpu().numpy()
formatted = (output * 255 / np.max(output)).astype("uint8")
depth = Image.fromarray(formatted)
depth = depth.convert('RGB')
return depth
if opts.control_move_processor:
self.model.to('cpu')
depth = Image.fromarray(formatted)
depth = depth.convert('RGB')
return depth