Upload 2 files

app.py

@@ -3,8 +3,8 @@ import image_feature as func
 
 
 def work11(image1, image2):
- return func.infer1(image1, image2)
- # return func.infer3(image1, image2)
+ # return func.infer1(image1, image2)
+ return func.infer3(image1, image2)
 
 
 # with gr.Blocks() as demo:

image_feature.py

@@ -49,16 +49,17 @@ DEVICE = torch.device('cpu')
 
 
 # 第二种方式推理图片相似度
-# processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
-# model = AutoModel.from_pretrained("google/vit-base-patch16-224").to(DEVICE)
-processor = AutoImageProcessor.from_pretrained("chanhua/autotrain-izefx-v3qh0")
-model = AutoModel.from_pretrained("chanhua/autotrain-izefx-v3qh0").to(DEVICE)
+processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
+model = AutoModel.from_pretrained("google/vit-base-patch16-224").to(DEVICE)
+# processor = AutoImageProcessor.from_pretrained("chanhua/autotrain-izefx-v3qh0")
+# model = AutoModel.from_pretrained("chanhua/autotrain-izefx-v3qh0").to(DEVICE)
 
 
 # tensor([0.6061], device='cuda:0', grad_fn=<SumBackward1>)
 
 # pipe = pipeline(task="image-feature-extraction", model_name="google/vit-base-patch16-384", device=DEVICE, pool=True)
-pipe = pipeline(task="image-feature-extraction", model_name="chanhua/autotrain-izefx-v3qh0", device=DEVICE, pool=True)
+# pipe = pipeline(task="image-feature-extraction", model_name="chanhua/autotrain-izefx-v3qh0", device=DEVICE, pool=True)
+pipe = pipeline(task="image-feature-extraction", model_name="google/vit-base-patch16-224", device=DEVICE, pool=True, revision="29e7a1e183")
 
 
 # 推理