import streamlit
from diffusers import StableDiffusionPipeline
import torch

model_id = "dream-textures/texture-diffusion"

if torch.cuda.is_available():
    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
    pipe = pipe.to("cuda")
    pipe.enable_xformers_memory_efficient_attention()
else:
    pipe = StableDiffusionPipeline.from_pretrained(model_id)

examplePrompt = "pbr brick wall"
streamlit.title("Texel Space Interpolation")

with streamlit.form(key="Interpolation"):
    prompt = streamlit.text_input(label="Prompt", value=examplePrompt)
    numSteps = streamlit.slider(label="Number of Inference steps", value=50, min_value=1, max_value=1000)
    numImages = streamlit.slider(label="Number of Generated Images", value=1, min_value=1, max_value=10)
    guidanceScale = streamlit.slider(label="Guidance Scale", value=7.5, min_value=0.0, max_value=100.0)
    images = pipe(prompt=prompt, num_inference_steps=numSteps, num_images_per_prompt=numImages, guidance_scale=guidanceScale).images
    streamlit.form_submit_button("Interpolate")
    synthesizedImage = streamlit.image(images)