import gradio as gr

import torch
import requests
from PIL import Image
from io import BytesIO

from diffusers import StableDiffusionImg2ImgPipeline

# load the pipeline
device = "cuda"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", revision="fp16", torch_dtype=torch.float16
).to(device)

# let's download an initial image
url = "https://stiltsdesigns.com/clients/pbm/hugging/mask.jpg"

response = requests.get(url)
init_image = Image.open(BytesIO(response.content)).convert("RGB")
init_image.thumbnail((768, 768))

# prompt = "A fantasy landscape, trending on artstation"

# images = pipe(prompt=prompt, image=init_image, strength=0.75, guidance_scale=7.5).images

# images[0].save("fantasy_landscape.png")

def spidey(prompt):
    images = pipe(prompt=prompt, image=init_image, strength=0.75, num_inference_steps=40, guidance_scale=10).images
    return images[0]

demo = gr.Interface(
    fn=spidey,
    inputs=gr.Textbox(lines=1, placeholder="Prompt..."),
    outputs="image",
)
demo.launch()