import gradio as gr


import PIL
import requests
import torch
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler

model_id = "timbrooks/instruct-pix2pix"
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16, safety_checker=None)
pipe.to("cpu")
pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)



def greet(name):
    prompt = "add tshirt"
    def download_image(img):
        image = PIL.Image.open(requests.get(img, stream=True).raw)
        image = PIL.ImageOps.exif_transpose(image)
        image = image.convert("RGB")
        return image
    image = download_image("https://raw.githubusercontent.com/timothybrooks/instruct-pix2pix/main/imgs/example.jpg")
    images = pipe(prompt, image=image, num_inference_steps=10, image_guidance_scale=1).images
      
    
    #image.save(images[0]+".png")
    return "done"

iface = gr.Interface(fn=greet, inputs="text", outputs="text")
iface.launch()