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]


iface = gr.Interface(fn=generate_view, inputs=[gr.Image(type="pil", value="images/bottle.png"), 
                                               gr.Number(label="elevation", value=0.), 
                                               gr.Number(label="azimuth", value=45.), 
                                               gr.Number(label="camera_distance", value=1.2), 
                                               gr.Number(label="num_inference_steps", value=20)], 
                                               outputs=gr.Image(width=256, height=256))
iface.launch()