from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, DPMSolverMultistepScheduler
import gradio as gr
import torch
from PIL import Image

model_id = 'Randolph/hadenjax-dreams'
prefix = ''
suffix = 'by hadenjax'
     
scheduler = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler")

pipe = StableDiffusionPipeline.from_pretrained(
  model_id,
  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
  scheduler=scheduler)

pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(
  model_id,
  torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
  scheduler=scheduler)

if torch.cuda.is_available():
  pipe = pipe.to("cuda")
  pipe_i2i = pipe_i2i.to("cuda")

def error_str(error, title="Error"):
    return f"""#### {title}
            {error}"""  if error else ""

def inference(prompt, guidance, steps, width=800, height=800, seed=0, img=None, strength=0.5, neg_prompt="", auto_prefix=False):

  generator = torch.Generator('cuda').manual_seed(seed) if seed != 0 else None
  prompt = f"{prefix} {prompt} {suffix}" if auto_prefix else f"{prompt} {suffix}"
  neg_prompt = f"{neg_prompt}, photo, DSLR, photorealistic"

  try:
    if img is not None:
      return img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height, generator), None
    else:
      return txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator), None
  except Exception as e:
    return None, error_str(e)

def txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator):

  result = pipe(
      prompt,
      negative_prompt = neg_prompt,
      num_inference_steps = int(steps),
      guidance_scale = guidance,
      width = width,
      height = height,
      generator = generator)
    
  return result.images[0]

def img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height, generator):

    ratio = min(height / img.height, width / img.width)
    img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)
    result = pipe_i2i(
        prompt,
        negative_prompt = neg_prompt,
        init_image = img,
        num_inference_steps = int(steps),
        strength = strength,
        guidance_scale = guidance,
        width = width,
        height = height,
        generator = generator)
        
    return result.images[0]

css = """.main-div div{display:inline-flex;align-items:center;gap:.8rem;font-size:1.75rem}.main-div div h1{font-weight:900;margin-bottom:7px}.main-div p{margin-bottom:10px;font-size:94%}a{text-decoration:underline}.tabs{margin-top:0;margin-bottom:0}#gallery{min-height:20rem}
"""
with gr.Blocks(css=css) as demo:
    gr.HTML(
        f"""<div class="main-div">
              <div>
                <h1>RIP Haden Jack Nimmer</h1>
                <h2>February 2nd, 1990 - May 31st, 2014</h2>
              </div>
              <p>
               <a href="https://huggingface.co/Randolph/hadenjax-dreams">HadenJax Dreams</a> is a memorial to my late brother, Haden "Jax" Jack Nimmer. It illustrates what you request in his artistic style.<br>
              </p>
              Running on {"<b>GPU 🔥</b>" if torch.cuda.is_available() else f"<b>CPU 🥶</b>. For faster inference it is recommended to <b>upgrade to GPU in <a href='https://huggingface.co/spaces/Randolph/hadenjax-dreams/settings'>Settings</a></b>"}<br><br>
            </div>
        """
    )   
    with gr.Row(equal_height=True):
        with gr.Column(scale=50):
            with gr.Group():
                image_out = gr.Image(height=800)
                error_output = gr.Markdown()
            with gr.Row():
                prompt = gr.Textbox(label="Prompt", show_label=False, max_lines=2, placeholder="(graphic novel)(webcomic of)(smudgy the whale)(parts-unknown)").style(container=False)
            with gr.Row():
                generate = gr.Button(value="Generate").style(rounded=(False, True, True, False))  
        with gr.Column(scale=50):
            with gr.Tab("Options"):
                with gr.Group():
                    neg_prompt = gr.Textbox(label="Negative prompt", placeholder="What to exclude from the image")
                    auto_prefix = gr.Checkbox(label="Unused", value=prefix, visible=prefix)
                with gr.Row():
                    guidance = gr.Slider(label="Guidance scale", value=10, maximum=15)
                with gr.Row():
                    steps = gr.Slider(label="Steps", value=25, minimum=2, maximum=75, step=1)
                with gr.Row():
                    width = gr.Slider(label="Width", value=400, minimum=64, maximum=1024, step=8)
                with gr.Row():
                    height = gr.Slider(label="Height", value=400, minimum=64, maximum=1024, step=8)
                with gr.Row():
                    seed = gr.Slider(0, 2147483647, label='Seed (0 = random)', value=0, step=1)
            with gr.Tab("Image to image"):
                with gr.Group():
                    image = gr.Image(label="Image", height=600, tool="editor", type="pil")
                    strength = gr.Slider(label="Transformation strength", minimum=0, maximum=1, step=0.01, value=0.5)

    auto_prefix.change(lambda x: gr.update(placeholder=f"[Your prompt] {suffix}"), inputs=auto_prefix, outputs=prompt, queue=False)
    inputs = [prompt, guidance, steps, width, height, seed, image, strength, neg_prompt, auto_prefix]
    outputs = [image_out, error_output]
    prompt.submit(inference, inputs=inputs, outputs=outputs)
    generate.click(inference, inputs=inputs, outputs=outputs)

    gr.HTML("""
    <div style="border-top: 1px solid #303030;">
      <br>
      <p>This space was created using <a href="https://huggingface.co/spaces/anzorq/sd-space-creator">SD Space Creator</a>.</p>
    </div>
    """)

demo.queue(concurrency_count=1)
demo.launch()