File size: 2,161 Bytes
ee3b2de
 
 
 
 
 
490262d
ee3b2de
 
 
 
 
 
 
af4ad56
ee3b2de
af4ad56
ee3b2de
af4ad56
 
ee3b2de
af4ad56
 
ee3b2de
af4ad56
 
e140357
ee3b2de
 
 
af4ad56
ee3b2de
 
af4ad56
ee3b2de
e83d19e
ee3b2de
af4ad56
 
ee3b2de
 
 
 
 
 
 
 
 
 
 
 
 
af4ad56
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
52
53
54

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
#print(f"SEED IS : {seed}")

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
    #print(f"seed is : {seed}")
    #print(f"image_in name is :{img_nm}")
    
    #if message == "revert": --to add revert functionality later
    if counter > 1:
      # Open the image
      image_in = Image.open("edited_image.png") #(img_nm) 
    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("edited_image.png") # (img_nm) #("./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(concurrency_count=10)
demo.launch(debug=True, width="80%", height=1500)