Browse Source

CLIPInterrogator -> Interrogator

replicate
pharmapsychotic 2 years ago
parent
commit
9ce0f68ab3
  1. 4
      clip_interrogator.ipynb
  2. 2
      clip_interrogator/__init__.py
  3. 2
      clip_interrogator/clip_interrogator.py
  4. 7
      run_cli.py

4
clip_interrogator.ipynb

@ -74,7 +74,7 @@
"\n", "\n",
"from PIL import Image\n", "from PIL import Image\n",
"\n", "\n",
"from clip_interrogator import CLIPInterrogator, Config\n" "from clip_interrogator import Interrogator, Config\n"
] ]
}, },
{ {
@ -91,7 +91,7 @@
"\n", "\n",
"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n", "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
"config = Config(device=device, clip_model_name=clip_model_name)\n", "config = Config(device=device, clip_model_name=clip_model_name)\n",
"interrogator = CLIPInterrogator(config)\n" "interrogator = Interrogator(config)\n"
] ]
}, },
{ {

2
clip_interrogator/__init__.py

@ -1,4 +1,4 @@
from .clip_interrogator import CLIPInterrogator, Config from .clip_interrogator import Interrogator, Config
__version__ = '0.1.3' __version__ = '0.1.3'
__author__ = 'pharmapsychotic' __author__ = 'pharmapsychotic'

2
clip_interrogator/clip_interrogator.py

@ -46,7 +46,7 @@ def _load_list(data_path, filename) -> List[str]:
return items return items
class CLIPInterrogator(): class Interrogator():
def __init__(self, config: Config): def __init__(self, config: Config):
self.config = config self.config = config
self.device = config.device self.device = config.device

7
main.py → run_cli.py

@ -7,7 +7,7 @@ import torch
from PIL import Image from PIL import Image
from clip_interrogator import CLIPInterrogator, Config from clip_interrogator import Interrogator, Config
def main(): def main():
@ -39,9 +39,8 @@ def main():
# generate a nice prompt # generate a nice prompt
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
config = Config(device=device, clip_model_name=args.clip) config = Config(device=device, clip_model_name=args.clip)
interrogator = CLIPInterrogator(config) ci = Interrogator(config)
prompt = interrogator.interrogate(image) print(ci.interrogate(image))
print(prompt)
if __name__ == "__main__": if __name__ == "__main__":
main() main()
Loading…
Cancel
Save