update omost (#190)

* update omost
This commit is contained in:
ZhouTianchen
2024-09-09 17:39:46 +08:00
committed by GitHub
parent 1887885274
commit 995f3374f1
5 changed files with 55 additions and 29 deletions

View File

@@ -55,11 +55,14 @@ class BasePipeline(torch.nn.Module):
def extend_prompt(self, prompt, local_prompts, masks, mask_scales):
local_prompts = local_prompts or []
masks = masks or []
mask_scales = mask_scales or []
extended_prompt_dict = self.prompter.extend_prompt(prompt)
prompt = extended_prompt_dict.get("prompt", prompt)
local_prompts += extended_prompt_dict.get("prompts", [])
masks += extended_prompt_dict.get("masks", [])
mask_scales += [5.0] * len(extended_prompt_dict.get("masks", []))
mask_scales += [100.0] * len(extended_prompt_dict.get("masks", []))
return prompt, local_prompts, masks, mask_scales
def enable_cpu_offload(self):