test.py

import gradio as gr
import torch
import numpy as np
from PIL import Image

# Define function to generate a red image
def generate_red_image(size: int=256):
    # Create a tensor with random values between 0 and 1 for each color channel
    image_tensor = torch.rand(3, size, size).to("cuda")

    image_tensor = image_tensor.to('cpu')

    # Convert Pytorch tensor to a numpy array and transpose axes to have the color channel last
    np_image = image_tensor.numpy().transpose(1, 2, 0)
    
    # Scale to 0-255 and convert to uint8
    np_image = (np_image * 255).astype(np.uint8)
    
    # Create an image from the numpy array
    img = Image.fromarray(np_image)
    
    return img

# Define the Gradio interface
iface = gr.Interface(fn=generate_red_image, inputs=gr.inputs.Slider(minimum=512, maximum=2048, step=64, default=512), outputs="image")
iface.queue(concurrency_count=64)

# Start the Gradio interface
iface.launch()