File size: 2,042 Bytes
ee3b2de 490262d ee3b2de e140357 ee3b2de e140357 ee3b2de e140357 ee3b2de e140357 ee3b2de e140357 ee3b2de e83d19e ee3b2de e140357 ee3b2de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 |
import PIL
import requests
import torch
import gradio as gr
import random
from PIL import Image
from diffusers import StableDiffusionInstructPix2PixPipeline, EulerAncestralDiscreteScheduler
model_id = "timbrooks/instruct-pix2pix"
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="fp16", safety_checker=None)
pipe.to("cuda")
pipe.enable_attention_slicing()
seed = random.randint(0, 1000000)
counter = 0
def chat(image_in, message, history): #, progress=gr.Progress(track_tqdm=True)):
#progress(0, desc="Starting...")
global counter
global seed
img_nm = f"./edited_image_{seed}.png"
counter += 1
#if message == "revert": --to add revert functionality later
if counter > 1:
# Open the image
image_in = Image.open(img_nm) #("./edited_image.png")
prompt = message #eg - "turn him into cyborg"
edited_image = pipe(prompt, image=image_in, num_inference_steps=20, image_guidance_scale=1).images[0]
edited_image.save(img_nm) #("edited_image.png") #("./edited_image.png")
history = history or []
add_text_list = ["There you go ", "Enjoy your image! ", "Nice work! Wonder what you gonna do next! ", "Way to go! ", "Does this work for you? ", "Something like this? "]
#Resizing the image for better display
response = random.choice(add_text_list) + '<img src="/file=' + img_nm[2:] + '" style="width: 200px; height: 200px;">'
#response = random.choice(add_text_list) + '<img src="/file=edited_image.png" style="width: 200px; height: 200px;">'
history.append((message, response))
return history, history
with gr.Blocks() as demo:
with gr.Row():
with gr.Column():
image_in = gr.Image(type='pil', label="Original Image")
text_in = gr.Textbox()
state_in = gr.State()
b1 = gr.Button('Edit the image!')
chatbot = gr.Chatbot()
b1.click(chat,[image_in, text_in, state_in], [chatbot, state_in])
#demo.queue()
demo.launch(debug=True, width="80%", height=1500) |