import gradio as gr
import torch
import requests
import json
# from diffusers import StableDiffusionLDM3DPipeline
import gradio as gr
import torch 
from PIL import Image
import base64
from io import BytesIO
from tempfile import NamedTemporaryFile
from pathlib import Path

from gradio_client import Client

theme = gr.themes.Base(
    primary_hue=gr.themes.Color(
        c100="#dbeafe", c200="#bfdbfe", c300="#93c5fd", c400="#60a5fa", c50="#eff6ff", c500="#0054ae", c600="#00377c", c700="#00377c", c800="#1e40af", c900="#1e3a8a", c950="#0a0c2b"),
    secondary_hue=gr.themes.Color(
        c100="#dbeafe", c200="#bfdbfe", c300="#93c5fd", c400="#60a5fa", c50="#eff6ff", c500="#0054ae", c600="#0054ae", c700="#0054ae", c800="#1e40af", c900="#1e3a8a", c950="#1d3660"),
).set(
    body_background_fill_dark='*primary_950',
    body_text_color_dark='#FFFFFF',
    body_text_color='#000000',
    border_color_accent='*primary_700',
    border_color_accent_dark='*neutral_800',
    block_background_fill_dark='*primary_950',
    block_border_width='2px',
    block_border_width_dark='2px',
    button_primary_background_fill_dark='*primary_500',
    button_primary_border_color_dark='*primary_500'
)

css='''
    @font-face {
        font-family: IntelOne;
        src: url("file/assets/intelone-bodytext-font-family-regular.ttf");
    }
'''

html_title = '''
<table>
<tr style="height:150px">
    <td style="border-bottom:0"><img src="file/assets/intel-labs.png" height="100" width="100"></td>
    <td style="border-bottom:0; vertical-align:middle">
    <p style="font-size:xx-large;font-family:IntelOne, Georgia, sans-serif;color: var(--body-text-color);">
     Cognitive AI:
     <br>
     LDM3D: Latent Diffusion Model for 3D
    </p>
    </td>
    <td style="border-bottom:0;"><img src="file/assets/gaudi.png" width="100" height="100"></td>
    <td style="border-bottom:0;"><img src="file/assets/xeon.png" width="100" height="100"></td>
</tr>
</table>

'''

client = Client("http://198.175.88.247:17810/")

def generate(
    prompt: str,
    negative_prompt: str,
    guidance_scale: float = 5.0,
    seed: int = 0,
    randomize_seed: bool = True,
):

    result = client.predict(
		prompt,	# str  in 'Prompt' Textbox component
		negative_prompt,	# str  in 'Negative Prompt' Textbox component
		guidance_scale,	# int | float (numeric value between 0 and 10) in 'Guidance Scale' Slider component
		seed,	# int | float (numeric value between 0 and 18446744073709551615) in 'Seed' Slider component
		randomize_seed,	# bool  in 'Randomize Seed' Checkbox component
		False,	# bool  in 'Upscale' Checkbox component
		fn_index=1
    )

    return result

with gr.Blocks(theme=theme, css=css) as demo:
    gr.HTML(value=html_title)
    gr.Markdown(
        """
[Model card](https://huggingface.co/Intel/ldm3d-pano<br>)
[Diffusers docs](https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/ldm3d_diffusion)
For better results, specify "360 view of" or "panoramic view of" in the prompt

"""
    )
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(label="Prompt")
            negative_prompt = gr.Textbox(label="Negative Prompt")
            guidance_scale = gr.Slider(
                label="Guidance Scale", minimum=0, maximum=10, step=0.1, value=5.0
            )
            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
            seed = gr.Slider(label="Seed", minimum=0,
                             maximum=2**64 - 1, step=1)
            generated_seed = gr.Number(label="Generated Seed")
            markdown = gr.Markdown(label="Output Box")
            with gr.Row():
                new_btn = gr.Button("New Image")
        with gr.Column(scale=2):
            html = gr.HTML()
            with gr.Row():
                rgb = gr.Image(label="RGB Image", type="filepath")
                depth = gr.Image(label="Depth Image", type="filepath")

    gr.Examples(
        examples=[
            ["360 view of a large bedroom", "", 7.0, 42, False]],
        inputs=[prompt, negative_prompt, guidance_scale, seed, randomize_seed],
        outputs=[rgb, depth, generated_seed, html],
        fn=generate,
        cache_examples=False)

    new_btn.click(
        fn=generate,
        inputs=[prompt, negative_prompt, guidance_scale, seed, randomize_seed],
        outputs=[rgb, depth, generated_seed, html],
    )
  
    demo.launch(
        allowed_paths=["assets/", "static/"]
    )