import torch
import requests
from PIL import Image
from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler
import rembg

# Load the pipeline
pipeline = DiffusionPipeline.from_pretrained(
    "sudo-ai/zero123plus-v1.1", custom_pipeline="sudo-ai/zero123plus-pipeline",
    torch_dtype=torch.float16
)

# Feel free to tune the scheduler!
# `timestep_spacing` parameter is not supported in older versions of `diffusers`
# so there may be performance degradations
# We recommend using `diffusers==0.20.2`
pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
    pipeline.scheduler.config, timestep_spacing='trailing'
)
pipeline.to('cuda:0')


def inference(input_img, num_inference_steps, guidance_scale, seed ):
    # Download an example image.
    cond = Image.open(input_img)
    
    # Run the pipeline!
    #result = pipeline(cond, num_inference_steps=75).images[0]
    result = pipeline(cond, num_inference_steps=num_inference_steps, 
                  guidance_scale=guidance_scale, 
                  generator=torch.Generator(pipeline.device).manual_seed(int(seed))).images[0]

    # for general real and synthetic images of general objects
    # usually it is enough to have around 28 inference steps
    # for images with delicate details like faces (real or anime)
    # you may need 75-100 steps for the details to construct
    
    #result.show()
    #result.save("output.png")
    return result

def remove_background(result):
    result = rembg.remove(result)
    return result


import gradio as gr

with gr.Blocks() as demo:
    gr.Markdown("<h1><center> Zero123++ Demo</center></h1>")
    with gr.Column():
      input_img = gr.Image(label='Input Image', type='filepath')
    with gr.Column():
      output_img = gr.Image(label='Zero123++ Output')
    with gr.Accordion("Advanced options:", open=False):
      rm_in_bkg = gr.Checkbox(label='Remove Input Background', )
      rm_out_bkg = gr.Checkbox(label='Remove Output Background')
      num_inference_steps = gr.Slider(label="Number of Inference Steps",minimum=15, maximum=100, step=1, value=75, interactive=True)
      guidance_scale = gr.Slider(label="Classifier Free Guidance Scale",minimum=1.00, maximum=10.00, step=0.1, value=4.0, interactive=True)
      seed = gr.Number(0, label='Seed')
    btn = gr.Button('Submit')

    btn.click(inference, [input_img, num_inference_steps, guidance_scale, seed ], output_img)
    rm_in_bkg.input(remove_background, input_img, output_img)
    rm_out_bkg.input(remove_background, output_img, output_img)
    
    gr.Examples(
    examples=[["extinguisher.png", 75, 4.0, 0],
              ['mushroom.png', 75, 4.0, 0], 
              ['tianw2.png', 75, 4.0, 0], 
              ['lysol.png', 75, 4.0, 0], 
              ['ghost-eating-burger.png', 75, 4.0, 0]
             ],
    inputs=[input_img, num_inference_steps, guidance_scale, seed],
    outputs=output_img,
    fn=inference,
    cache_examples=True,
      )


demo.launch()