sergiopaniego's picture
Updated app
f0585ee
raw
history blame
621 Bytes
import gradio as gr
import spaces
import torch
from PIL import Image
import requests
from transformers import DetrImageProcessor
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
zero = torch.Tensor([0]).cuda()
print(zero.device) # <-- 'cpu' πŸ€”
@spaces.GPU
def detect(image):
encoding = processor(image, return_tensors='pt')
print(encoding.keys())
print(zero.device) # <-- 'cuda:0' πŸ€—
return f"Hello {encoding.keys()} Tensor"
demo = gr.Interface(fn=detect, inputs=gr.Image(label="Input image", type="pil"), outputs=gr.Image(label="Output image", type="pil"))
demo.launch()