File size: 2,032 Bytes
c04f31b 78f4578 d997936 c04f31b 17b4c6b 14d6567 17b4c6b c04f31b 78f4578 c04f31b 0858412 78f4578 c04f31b 0858412 78f4578 c04f31b f52d753 78f4578 1752ad5 78f4578 1752ad5 78f4578 1752ad5 78f4578 1752ad5 c04f31b f52d753 |
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 |
import gradio as gr
from transformers import AutoTokenizer, VisionEncoderDecoderModel, ViTImageProcessor
from diffusers import StableDiffusionPipeline
# Initialize device and models for captioning
device = 'cpu'
encoder_checkpoint = "nlpconnect/vit-gpt2-image-captioning"
decoder_checkpoint = "nlpconnect/vit-gpt2-image-captioning"
model_checkpoint = "nlpconnect/vit-gpt2-image-captioning"
feature_extractor = ViTImageProcessor.from_pretrained(encoder_checkpoint)
tokenizer = AutoTokenizer.from_pretrained(decoder_checkpoint)
caption_model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint).to(device)
# Load the Stable Diffusion model
diffusion_model = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
diffusion_model = diffusion_model.to(device)
def get_caption(image):
# Generate a caption from the image
image = image.convert('RGB')
image_tensor = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device)
caption_ids = caption_model.generate(image_tensor, max_length=128, num_beams=3)[0]
caption_text = tokenizer.decode(caption_ids, skip_special_tokens=True)
return caption_text
def generate_image(caption):
# Generate an image from the caption
generated_image = diffusion_model(caption)["sample"][0]
return generated_image
# Set up Gradio interface
title = "Image Captioning and Generation"
with gr.Blocks(title=title) as demo:
with gr.Row():
with gr.Column():
image_input = gr.Image(label="Upload any Image", type='pil')
get_caption_btn = gr.Button("Get Caption")
with gr.Column():
caption_output = gr.Textbox(label="Caption")
generate_image_btn = gr.Button("Generate Image")
with gr.Row():
generated_image_output = gr.Image(label="Generated Image")
get_caption_btn.click(get_caption, inputs=image_input, outputs=caption_output)
generate_image_btn.click(generate_image, inputs=caption_output, outputs=generated_image_output)
demo.launch()
|