import spaces 
import gradio as gr 
import platform
import os;
import socket;
import torch
import torch.nn as nn

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()

    def forward(self, x):
        pass
        
    def device():
       return next(model.parameters()).device

model = SimpleModel().to(device);

@torch.inference_mode()
def sysinfo():
    RandomTensor = torch.randn(1, 2)  # Example audio tensor
    tensorExample = RandomTensor.to(device)
    
    return f"""
       hostname: {platform.node()} {socket.gethostname()}
       device: {device}
       model device: {model.device}
       tensor: {tensorExample}
    """;

@spaces.GPU
def gpu():
     return sysinfo();


def nogpu():

    return sysinfo();



with gr.Blocks() as demo:
    outgpu = gr.Textbox(lines=5);
    outnpu = gr.Textbox(lines=5);
    btngpu = gr.Button(value="gpu");
    btngpun = gr.Button(value="ngpu");
    
    btngpu.click(gpu, None, [outgpu]);
    btngpun.click(nogpu, None, [outnpu]);
    


if __name__ == "__main__":
    demo.launch(
        share=False,
        debug=False,
        server_port=7860,
        server_name="0.0.0.0"
    )