import os
import json
import random
from getpass import getpass
from glob import glob
from natsort import natsorted
import matplotlib.pyplot as plt
from PIL import Image

import torch
from torch import autocast
from diffusers import StableDiffusionPipeline, DDIMScheduler
from IPython.display import display


import gradio as gr

from gradio.components import Textbox, Image


repo_name = 'mohansathya/twosd' # YOUR REPO NAME
pipe2 = StableDiffusionPipeline.from_pretrained(repo_name, torch_dtype=torch.float16).to('cuda')



def generate_query_response(prompt):
    negative_prompt = "bad anatomy, ugly, deformed, desfigured, distorted, poorly drawn, blurry, low quality, low definition, lowres, out of frame, out of image, cropped, cut off, signature, watermark"
    num_samples = 5
    guidance_scale = 7.5
    num_inference_steps = 30
    height = 512
    width = 512

    seed = random.randint(0, 2147483647)
    print("Seed: {}".format(str(seed)))
    generator = torch.Generator(device='cuda').manual_seed(seed)

    with autocast("cuda"), torch.inference_mode():
        imgs = pipe2(
            prompt,
            negative_prompt=negative_prompt,
            height=height, width=width,
            num_images_per_prompt=num_samples,
            num_inference_steps=num_inference_steps,
            guidance_scale=guidance_scale,
            generator=generator
        ).images

    for img in imgs:
        return img

    
# Input from user
in_prompt = Textbox(label="Enter a prompt:")

# Output response
out_response = Image(label="Generated image:")
    
# Gradio interface to generate UI link
iface = gr.Interface(
    fn=generate_query_response, inputs=in_prompt, outputs=out_response) 

# Launch the interface to generate UI
iface.launch()