hjxdl 0.1.55__py3-none-any.whl → 0.1.56__py3-none-any.whl

hdl/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.1.55'
-__version_tuple__ = version_tuple = (0, 1, 55)
+__version__ = version = '0.1.56'
+__version_tuple__ = version_tuple = (0, 1, 56)

hdl/utils/llm/vis.py

@@ -21,13 +21,22 @@ class ImgHandler:
         model_path,
         redis_host,
         redis_port,
+        device: str = None
     ) -> None:
-        self.model = ChineseCLIPModel.from_pretrained(model_path)
+        if device is None:
+            self.device = torch.device("cuda") \
+                if torch.cuda.is_available() \
+                else torch.device("cpu")
+        else:
+            self.device = device
+
+        self.model = ChineseCLIPModel.from_pretrained(model_path).to(self.device)
         self.processor = ChineseCLIPProcessor.from_pretrained(model_path)
         self.redis_host = redis_host
         self.redis_port = redis_port
         self._redis_conn = None
 
+
     @property
     def redis_conn(self):
         if self._redis_conn is None:
@@ -39,7 +48,7 @@ class ImgHandler:
             images=images,
             return_tensors="pt",
             **kwargs
-        )
+        ).to(self.device)
         image_features = self.model.get_image_features(**inputs)
         image_features = image_features / \
             image_features.norm(p=2, dim=-1, keepdim=True)
@@ -55,7 +64,7 @@ class ImgHandler:
             padding=True,
             return_tensors="pt",
             **kwargs
-        )
+        ).to(self.device)
         text_features = self.model.get_text_features(**inputs)
         text_features = text_features / \
             text_features.norm(p=2, dim=-1, keepdim=True)

{hjxdl-0.1.55.dist-info → hjxdl-0.1.56.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: hjxdl
-Version: 0.1.55
+Version: 0.1.56
 Summary: A collection of functions for Jupyter notebooks
 Home-page: https://github.com/huluxiaohuowa/hdl
 Author: Jianxing Hu

{hjxdl-0.1.55.dist-info → hjxdl-0.1.56.dist-info}/RECORD

@@ -1,5 +1,5 @@
 hdl/__init__.py,sha256=GffnD0jLJdhkd-vo989v40N90sQbofkayRBwxc6TVhQ,72
-hdl/_version.py,sha256=iWADRf4n340eBsELvrfW_j6SevUqYhLDta9JVrRjHMo,413
+hdl/_version.py,sha256=oYxzoSIQhFUGS4XB8_o2EJqTMxHBvofvy2vOYfOC7c0,413
 hdl/args/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hdl/args/loss_args.py,sha256=s7YzSdd7IjD24rZvvOrxLLFqMZQb9YylxKeyelSdrTk,70
 hdl/controllers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -131,12 +131,12 @@ hdl/utils/llm/chat.py,sha256=sk7Lw5Oa30k-l2fnJknkMmTc5zkBeEKsR981aeFhH5s,11907
 hdl/utils/llm/embs.py,sha256=Tf0FOYrOFZp7qQpEPiSCXzlgyHH0X9HVTUtsup74a9E,7174
 hdl/utils/llm/extract.py,sha256=2sK_WJzmYIc8iuWaM9DA6Nw3_6q1O4lJ5pKpcZo-bBA,6512
 hdl/utils/llm/llama_chat.py,sha256=watcHGOaz-bv3x-yDucYlGk5f8FiqfFhwWogrl334fk,4387
-hdl/utils/llm/vis.py,sha256=l98iSAAUeANBhzhGbho5pXzGhMS73SqIqXgcUtLciBI,2176
+hdl/utils/llm/vis.py,sha256=xGHtLlCYHFZHQQzwA3eq3U0T7WauyV1-qWmpE5ds3yg,2464
 hdl/utils/schedulers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hdl/utils/schedulers/norm_lr.py,sha256=bDwCmdEK-WkgxQMFBiMuchv8Mm7C0-GZJ6usm-PQk14,4461
 hdl/utils/weather/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 hdl/utils/weather/weather.py,sha256=k11o6wM15kF8b9NMlEfrg68ak-SfSYLN3nOOflFUv-I,4381
-hjxdl-0.1.55.dist-info/METADATA,sha256=2pTlv2FRyTOZ7aRwIqn5kfXdRIP9-8BOQTOhhTfZMk0,849
-hjxdl-0.1.55.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
-hjxdl-0.1.55.dist-info/top_level.txt,sha256=-kxwTM5JPhylp06z3zAVO3w6_h7wtBfBo2zgM6YZoTk,4
-hjxdl-0.1.55.dist-info/RECORD,,
+hjxdl-0.1.56.dist-info/METADATA,sha256=ObGUeF2jiHrJk7Wizw8_1axb0t1KxPvWkpzPVHYOmX4,849
+hjxdl-0.1.56.dist-info/WHEEL,sha256=cVxcB9AmuTcXqmwrtPhNK88dr7IR_b6qagTj0UvIEbY,91
+hjxdl-0.1.56.dist-info/top_level.txt,sha256=-kxwTM5JPhylp06z3zAVO3w6_h7wtBfBo2zgM6YZoTk,4
+hjxdl-0.1.56.dist-info/RECORD,,