Update app.py

app.py

@@ -1,65 +1,65 @@
-import gradio as gr
-import numpy as np
-import random
-import time
-import json
-import os
-from loguru import logger
-from decouple import config
-import io
-import torch
-import numpy as np
-import torch
-import cv2
-from PIL import Image
-
-from segment_anything import sam_model_registry, SamPredictor
-
-import spaces
-
-print(f"Is CUDA available: {torch.cuda.is_available()}")
-print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
-print(torch.version.cuda)
-device = torch.cuda.get_device_name(torch.cuda.current_device())
-print(device)
-
-sam_checkpoint = "sam-hq/models/sam_hq_vit_h.pth"
-model_type = "vit_h"
-device = "cuda"
-sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
-sam.to(device=device)
-predictor = SamPredictor(sam)
-
-@spaces.GPU(duration=5)
-def generate_image(prompt, image):
-    predictor.set_image(image)
-
-    prompt = json.loads(prompt)
-    input_points = np.array(prompt['input_points'])
-    input_labels = np.array(prompt['input_labels'])
-
-    mask, _, _ = predictor.predict(
-        point_coords=input_points,
-        point_labels=input_labels,
-        box=None,
-        multimask_output=False,
-        hq_token_only=True,
-    )
-
-    rgb_array = np.zeros((mask.shape[1], mask.shape[2], 3), dtype=np.uint8)
-    rgb_array[mask[0]] = 255
-    result = Image.fromarray(rgb_array)
-
-    return result
-
-
-if __name__ == "__main__":
-    demo = gr.Interface(fn=generate_image, inputs=[
-        "text",
-        gr.Image(image_mode='RGB', type="numpy")
-    ],
-                        outputs=[
-                            gr.Image(type="numpy", image_mode='RGB')
-                        ])
-    demo.launch(debug=True)
-    logger.debug('demo.launch()')
+import gradio as gr
+import numpy as np
+import random
+import time
+import json
+import os
+from loguru import logger
+from decouple import config
+import io
+import torch
+import numpy as np
+import torch
+import cv2
+from PIL import Image
+
+from segment_anything_hq import sam_model_registry, SamPredictor
+
+import spaces
+
+print(f"Is CUDA available: {torch.cuda.is_available()}")
+print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+print(torch.version.cuda)
+device = torch.cuda.get_device_name(torch.cuda.current_device())
+print(device)
+
+sam_checkpoint = "sam_hq_vit_h.pth"
+model_type = "vit_h"
+device = "cuda"
+sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
+sam.to(device=device)
+predictor = SamPredictor(sam)
+
+@spaces.GPU(duration=5)
+def generate_image(prompt, image):
+    predictor.set_image(image)
+
+    prompt = json.loads(prompt)
+    input_points = np.array(prompt['input_points'])
+    input_labels = np.array(prompt['input_labels'])
+
+    mask, _, _ = predictor.predict(
+        point_coords=input_points,
+        point_labels=input_labels,
+        box=None,
+        multimask_output=False,
+        hq_token_only=True,
+    )
+
+    rgb_array = np.zeros((mask.shape[1], mask.shape[2], 3), dtype=np.uint8)
+    rgb_array[mask[0]] = 255
+    result = Image.fromarray(rgb_array)
+
+    return result
+
+
+if __name__ == "__main__":
+    demo = gr.Interface(fn=generate_image, inputs=[
+        "text",
+        gr.Image(image_mode='RGB', type="numpy")
+    ],
+                        outputs=[
+                            gr.Image(type="numpy", image_mode='RGB')
+                        ])
+    demo.launch(debug=True)
+    logger.debug('demo.launch()')