import gradio as gr
import torch
import diffusers
from diffusers import DiffusionPipeline
from zero123 import Zero123Pipeline
diffusers.Zero123Pipeline = Zero123Pipeline

def generate_view(source_img, elevation, azimuth, camera_distance, num_inference_steps):
        
    if torch.cuda.is_available():
        device = 'cuda:0'
    else:
        device = 'cpu'

    # Prepare pipeline
    pipeline = DiffusionPipeline.from_pretrained("ashawkey/stable-zero123-diffusers", trust_remote_code=True)
    pipeline.to(device)

    # Prepare input data:
    image = source_img.resize((256, 256)).convert("RGB")


    # Generate and save images:
    images = pipeline([image],
                    torch.tensor([elevation], dtype=torch.float16).to(device),
                    torch.tensor([azimuth], dtype=torch.float16).to(device),
                    torch.tensor([camera_distance], dtype=torch.float16).to(device),
                    num_inference_steps=int(num_inference_steps)).images

    return images[0]


def app():
    with gr.Blocks():
        with gr.Row():
            with gr.Column():
                image = gr.Image(type="pil")

                elevation = gr.Number(label="elevation", value=0.)
                azimuth = gr.Number(label="azimuth", value=45.) 
                camera_distance = gr.Number(label="camera_distance", value=1.2)
                num_inference_steps = gr.Slider(
                    label="Inference Steps",
                    minimum=0,
                    maximum=100,
                    step=1,
                    value=20,
                )

                generate = gr.Button(value="Generate")

            with gr.Column():
                output = gr.Image(type="pil",label="Output", width=256, height=256)

        generate.click(
            fn=generate_view,
            inputs=[
                image,
                elevation,
                azimuth,
                camera_distance,
                num_inference_steps,
            ],
            outputs=[output],
        )
        
        gr.Examples(
            examples=[
                [
                    "images/bottle.png",
                    0.,
                    45.,
                    1.2,
                    40,
                ]
            ],
            fn=generate_view,
            inputs=[
                image,
                elevation,
                azimuth,
                camera_distance,
                num_inference_steps,
            ],
            outputs=[output],
            cache_examples=True,
        )


gradio_app = gr.Blocks()
with gradio_app:
    gr.HTML(
        """
    <h1 style='text-align: center'>
    Demo of new view generation with <a href='https://huggingface.co/stabilityai/stable-zero123' target='_blank'>Stable-Zero123</a> implemented in <a href='https://huggingface.co/ashawkey/stable-zero123-diffusers' target='_blank'>stable-zero123-diffusers</a>
    </h1>
    """)
    with gr.Row():
        with gr.Column():
            app()

gradio_app.launch(debug=True)