switch to api config

app.py

@@ -11,49 +11,48 @@ import fastai
 from deoldify import device
 from deoldify.device_id import DeviceId
 from deoldify.visualize import *
-from huggingface_hub import snapshot_download
+# from huggingface_hub import HfApi, snapshot_download
+from huggingface_hub import HfApi
 
 os.system("pip freeze")
 from collections.abc import Sized # Import Sized from collections.abc
+# Suppress warnings
 warnings.filterwarnings("ignore", category=UserWarning, message=".*?Your .*? set is empty.*?")
 
+# Initialize Hugging Face API
+api = HfApi()
 
-# private repo
+device.set(device=DeviceId.CPU)
+
+# Private repo settings
 token = os.getenv("HF_TOKEN")
 repo_id = "afondiel/image-colorizer-deoldify"
 repo_type = "space"
 
 # Download the snapshot from the space repository
-snapshot_folder = snapshot_download(repo_id=repo_id, repo_type=repo_type, token=token)
-
-# Set the device to use for computation
-# choices:  CPU, GPU0...GPU7
-if torch.cuda.is_available():
-    device.set(device=DeviceId.GPU0)
-else:
-    device.set(device=DeviceId.CPU)
+snapshot_folder = api.snapshot_download(repo_id=repo_id, repo_type=repo_type, token=token)
 
+device.set(device=DeviceId.GPU0)
 # Load the pre-trained model
 _colorizer = get_image_colorizer(root_folder=Path(snapshot_folder), artistic=True)
 
-
 def colorizer_fn(input_img, render_factor):
-  """
-  Colorize grayscale images/photos
-  - @param input_img old (grayscale) image 
-  - @param render_factor render_factor
-  """
-  if input_img is not None and input_img !='':
-    output_img = _colorizer.get_transformed_image(
-        path=input_img,
-        render_factor=int(render_factor),
-        watermarked=watermarked,
-        post_process=True,
-    )
-  else:
-    print('Provide an image and try again.')
+    """
+    Colorize grayscale images/photos
+    - @param input_img old (grayscale) image 
+    - @param render_factor render_factor
+    """
+    if input_img is not None and input_img != '':
+        output_img = _colorizer.get_transformed_image(
+            path=input_img,
+            render_factor=int(render_factor),
+            watermarked=watermarked,
+            post_process=True,
+        )
+    else:
+        print('Provide an image and try again.')
 
-  return (input_img, output_img) # Return a tuple of old and color Image to be plotted with ImageSlider()
+    return (input_img, output_img)  # Return a tuple of old and color Image to be plotted with ImageSlider()
 
 title = "AI Image Colorizer"
 description = "Colorize old images with AI"
@@ -61,7 +60,7 @@ examples = [["./demo.jpg"],]
 
 demo = gr.Interface(
     fn=colorizer_fn,
-    inputs=[gr.Image(type="filepath" , label="Old image"), gr.Slider(0, 40, label="Render Factor", value=10)],
+    inputs=[gr.Image(type="filepath", label="Old image"), gr.Slider(0, 40, label="Render Factor", value=10)],
     outputs=ImageSlider(type="pil", label="Old vs Colored image"),
     examples=examples,
     title=title,
@@ -69,6 +68,5 @@ demo = gr.Interface(
 )
 
 # Launch the demo
-
 if __name__ == "__main__":
-  demo.launch()
+    demo.launch()