from googletrans import Translator
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import torch
import gradio as gr

translator = Translator()

model_id = "stabilityai/stable-diffusion-2-1"
access_token="your token"

pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16, use_auth_token=access_token
)

pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

def generate(prompt, inference_steps, guidance_scale, neg_prompt):
  if not neg_prompt:
    neg_prompt = ""
    
    prompt_eng = translator.translate(prompt, dest='en').text
    image = pipe(prompt_eng, guidance_scale= int(guidance_scale), num_inference_steps = int(inference_steps), negative_prompt = neg_prompt).images[0]
    return image

gr.Interface(
    generate,
    title = 'Image to Image using Diffusers',
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Slider(50, 700, value=50, label ="Inference steps"),
        gr.Slider(1, 10, value=5, label ="Guidance scale"),
        gr.Textbox(label="Negative prompt (include things you DO NOT want in the image")
    ],
    outputs = [
        gr.Image(elem_id="output-image"),

        ], css = "#output-image, #input-image, #image-preview {border-radius: 40px !important; background-color : gray !important;} "
).launch()