import torch
from diffusers import StableDiffusionPipeline

model_id = "runwayml/stable-diffusion-v1-5"
sd_pipeline = StableDiffusionPipeline.from_pretrained(model_id,
                                                      torch_dtype=torch.float16)
sd_pipeline = sd_pipeline.to("cuda")

def get_completion_sd(prompt):
  negative_prompt = """
  simple background, duplicate, low quality, lowest quality,
  bad anatomy, bad proportions, extra digits, lowres, username,
  artist name, error, duplicate, watermark, signature, text,
  extra digit, fewer digits, worst quality, jpeg artifacts, blurry
  """
  return sd_pipeline(prompt, negative_prompt=negative_prompt).images[0]



#let's prompt

# prompt = "astronaut, riding a horse, on mars, human colony"
# prompt = "children, playing in disneyland, view from a distance"
prompt = """llama, wearing red socks,
grazing, open field, raining
"""
print(prompt)
sd_image = get_completion_sd(prompt)
sd_image.save("./llama.jpg")



import gradio as gr

def get_completion(prompt):
  negative_prompt = """
  simple background, duplicate, low quality, lowest quality,
  bad anatomy, bad proportions, extra digits, lowres, username,
  artist name, error, duplicate, watermark, signature, text,
  extra digit, fewer digits, worst quality, jpeg artifacts, blurry
  """
  return sd_pipeline(prompt, negative_prompt=negative_prompt).images[0]

# def generate(prompt):
#     output = get_completion_sd(prompt)
#     return output

genai_app = gr.Interface(fn=get_completion,
                         inputs=[gr.Textbox(label="Your prompt")],
                         outputs=[gr.Image(label="Result")],
                         title="Generate Cool Images",
                         description="Generate any image with Stable Diffusion",
                         allow_flagging="never",
                         examples=["astronaut, riding a horse, on mars",
                                   "cargo ship, flying, in space"])
genai_app.launch(share=True)