Spaces:
Sleeping
Sleeping
File size: 1,125 Bytes
8d64162 70d9de4 8d64162 0d77bb1 d197e7f 96bca50 0d77bb1 96bca50 d197e7f 96bca50 d197e7f 70d9de4 af715dd 70d9de4 8d64162 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 |
import base64
import io
import torch
from PIL import Image
import wandb
def get_wandb_artifact(
artifact_name: str,
artifact_type: str,
get_metadata: bool = False,
) -> str:
if wandb.run:
artifact = wandb.use_artifact(artifact_name, type=artifact_type)
artifact_dir = artifact.download()
else:
api = wandb.Api()
artifact = api.artifact(artifact_name)
artifact_dir = artifact.download()
if get_metadata:
return artifact_dir, artifact.metadata
return artifact_dir
def get_torch_backend():
if torch.cuda.is_available():
if torch.backends.cuda.is_built():
return "cuda"
if torch.backends.mps.is_available():
if torch.backends.mps.is_built():
return "mps"
return "cpu"
return "cpu"
def base64_encode_image(image: Image.Image, mimetype: str) -> str:
byte_arr = io.BytesIO()
image.save(byte_arr, format="PNG")
encoded_string = base64.b64encode(byte_arr.getvalue()).decode("utf-8")
encoded_string = f"data:{mimetype};base64,{encoded_string}"
return str(encoded_string)
|