import gradio as gr
import spaces
import torch
from huggingface_hub import InferenceClient

zero = torch.Tensor([0]).cuda()

print(zero.device) # <-- 'cpu' 🤔

client = InferenceClient()
image_model = 'black-forest-labs/FLUX.1-dev'

@spaces.GPU
def generate_image(input_prompt):
    image = client.text_to_image(prompt=input_prompt, model=image_model)
    return(image)

with gr.Blocks(title="ball") as demo:

    input_prompt = gr.Textbox(label="Prompt")

    text_to_image_generate_button = gr.Button("Generate")

    text_to_image_output = gr.Image(label="Image", type="pil")

    text_to_image_generate_button.click(
        fn=generate_image,
        inputs=input_prompt,
        outputs=text_to_image_output
    )
    demo.launch()