Spaces:
Runtime error
Runtime error
| 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) | |