import os
import gradio as gr
import PIL.Image
import numpy as np
import random
import torch
import subprocess
from diffusers import StableDiffusionPipeline
# from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import time

model_id = "dicoo_model"

# dpm = DPMSolverMultistepScheduler.from_pretrained(model_id, subfolder="scheduler")
# pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=dpm, torch_dtype=torch.float)
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float)

def predict(prompt, steps=50, seed=42, guidance_scale=7.5):
    print("prompt: ", prompt)
    print("steps: ", steps)
    generator = torch.manual_seed(seed)
    start_time = time.time()
    image = pipe(prompt, generator=generator, num_inference_steps=steps, guidance_scale=7.5).images[0]
    print("cost: ", time.time() - start_time)

    return image

md = """
This Spaces app is created by Intel AIA/AIPC team with the model fine-tuned with one shot (one image) for a newly introduced object \"dicoo\". To replicate the model fine-tuning, please refer to the code sample in <a href=\"https://github.com/intel/neural-compressor/tree/master/examples/pytorch/diffusion_model/diffusers/textual_inversion\">Intel Neural Compressor</a>. You may also refer to our <a href=\"https://medium.com/intel-analytics-software/personalized-stable-diffusion-with-few-shot-fine-tuning-on-a-single-cpu-f01a3316b13\">blog</a> for more details. **For better experience and faster generation, please visit <a href=\"https://huggingface.co/spaces/lvkaokao/INC-Dicoo-Diffusion\">this page</a>, with <a href=\"https://arxiv.org/abs/2206.00927\">DPM-Solver scheduler</a> and upgraded hardware support.**

You can skip the queue by duplicating this space: <a style="display:inline-block" href="https://huggingface.co/spaces/Intel/dicoo_diffusion?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a>"""

random_seed = random.randint(0, 2147483647)
gr.Interface(
    predict,
    inputs=[
        gr.inputs.Textbox(label='Prompt', default='a lovely <dicoo> in red dress and hat, in the snowy and brightly night, with many brightly buildings'),
        gr.inputs.Slider(1, 100, label='Inference Steps', default=50, step=1),
        gr.inputs.Slider(0, 2147483647, label='Seed', default=random_seed, step=1),
        gr.inputs.Slider(1.0, 20.0, label='Guidance Scale - how much the prompt will influence the results', default=6.0, step=0.1),
    ],
    outputs=gr.Image(shape=[512, 512], type="pil", elem_id="output_image"),
    css="#output_image{width: 256px}",
    title="Demo of dicoo-finetuned-diffusion-model 🧨",
    description=md,
).launch()