import gradio as gr
import torch
from gradio.components import Image
from diffusers import StableDiffusionPipeline
from huggingface_hub import login
import os
from dotenv import load_dotenv

load_dotenv()

login(token=os.getenv('HF_W_TOKEN'))

name_of_your_concept = "sks"  # CHANGE THIS ACCORDING TO YOUR SUBJECT
type_of_thing = "boy"  # CHANGE THIS ACCORDING TO YOUR SUBJECT

# Tune the guidance to control how closely the generations follow the prompt.
# Values between 7-11 usually work best
guidance_scale = 11


with torch.no_grad():
    torch.cuda.empty_cache()

pipe = StableDiffusionPipeline.from_pretrained(
    'raulm/hebru-small-batch-dreambooth',
    torch_dtype=torch.float16,
).to("cuda")


def generation_fn(prompt):
    print(prompt)
    imgs = pipe(prompt=prompt, guidance_scale=guidance_scale).images
    return imgs[0]


gr.Interface(generation_fn,
             "text",
             Image(type="pil", label="Generated image"),
             ).launch()