import os
import json
import random

import torch
from torch import autocast
from diffusers import StableDiffusionPipeline, DDIMScheduler


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.bfloat16)



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 = 6
    height = 512
    width = 512

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

    with autocast("cpu", dtype=torch.bfloat16), 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()