adding image gen

app.py

@@ -1,14 +1,31 @@
 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 greet(n):
-    print(zero.device) # <-- 'cuda:0' 🤗
-    return f"Hello {zero + n} Tensor"
+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")
 
-demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
-demo.launch()
+    text_to_image_generate_button.click(
+        fn=generate_image,
+        inputs=input_prompt,
+        outputs=text_to_image_output
+    )
+    demo.launch()