import gradio as gr
import torch
from utils import load_model, generate_random_img, schedule_function
import time
import random
import threading 
from gradio_client import Client

def generate_image():
    with torch.no_grad():
        model = load_model('generator', 'generator_model_epoch_94.pth')
        generated_image = generate_random_img(model)
        return generated_image

iface = gr.Interface(
    fn=generate_image,
    inputs=[],
    outputs=gr.outputs.Image(type='numpy'),
    allow_screenshot=True,
    title='Random Landscape Image Generator By Huseyn Gorbani',
    description='This app generates random images, using DCFAN inspired WGAN-GP model. Special Thanks to Aladdin Persson and Emilien Dupont for their insightful repos on GitHub. Aladdin Persson (repo: https://github.com/aladdinpersson/Machine-Learning-Collection/tree/master/ML/Pytorch/GANs/4.%20WGAN-GP) Emilien Dupont (repo: https://github.com/EmilienDupont/wgan-gp/blob/master/training.py)',
    css='img_styles.css',
)

if __name__ == '__main__':
    scheduler_thread = threading.Thread(target=schedule_function) # avoiding sleep, again this project is for academic purposes only
    # scheduler_thread.start()
    iface.launch()