Spaces:
Runtime error
Runtime error
File size: 3,557 Bytes
0bfa5c3 | 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import gradio as gr
from PIL import Image
from loguru import logger
import torch
from diffusers import AutoPipelineForImage2Image
from diffusers import AutoPipelineForText2Image
# Init pipelines
pipe_t2i = AutoPipelineForText2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
pipe_t2i.to("cuda")
pipe_i2i = AutoPipelineForImage2Image.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16")
pipe_i2i.to("cuda")
def generate_image(text_input, prev_image, strength):
if prev_image:
logger.info(f"iterating for prompt: {text_input} with strength {strength}")
image = pipe_i2i(
prompt=text_input,
num_inference_steps=4,
guidance_scale=0.0,
strength=strength,
image=prev_image
).images[0]
else:
logger.info(f"generating for prompt: {text_input}")
image = pipe_t2i(text_input, num_inference_steps=4, guidance_scale=0.0).images[0]
return image
# Create left block (Chatbot)
with gr.Blocks() as demo:
with gr.Row() as row:
with gr.Column():
# create left block (Text input)
chatbot = gr.Chatbot(label="History")
msg = gr.Textbox(label="Input prompt and press Enter")
strength_slider = gr.Slider(minimum=0, maximum=1, value=0.75, label="Creativity of image iterations")
regen = gr.Button("Regenerate for the last prompt")
clear = gr.Button("Clear")
def respond(message, chat_history, gen_image, strength_slider, prev_gen_image):
user_input = [chat_history[i][0] for i in range(len(chat_history))]
user_input.append(message)
bot_message = "Generating image for: " + " ".join(user_input)
chat_history.append((message, bot_message))
if gen_image:
prev_gen_image = gen_image
else:
prev_gen_image = None
gen_image = generate_image(" ".join(user_input), gen_image, strength_slider)
return "", chat_history, gen_image, prev_gen_image
def regenerate(message, chat_history, gen_image, strength_slider):
# Regenerate with the previous image
user_input = [chat_history[i][0] for i in range(len(chat_history))]
# user_input.append(message)
bot_message = "Regenerating image for: " + " ".join(user_input)
chat_history.append((message, bot_message))
gen_image = generate_image(" ".join(user_input), gen_image, strength_slider)
return "", chat_history, gen_image
with gr.Column():
# Create right block (Image Generator)
image_generator = gr.Image(label="Generated image", type="pil", interactive=False, show_download_button=True)
# use this to keep track of previous image generation but don't display it
image_generator_prev = gr.Image(label="Previous image", type="pil", interactive=False)
msg.submit(respond, [msg, chatbot, image_generator, strength_slider, image_generator_prev], [msg, chatbot, image_generator, image_generator_prev])
regen.click(regenerate, [msg, chatbot, image_generator_prev, strength_slider], [msg, chatbot, image_generator])
clear.click(lambda: [None, None, None, 0.75], None, [chatbot, image_generator, image_generator_prev, strength_slider], queue=False)
# Launch the Gradio app
demo.launch(share=True)
|