import torch
import gradio as gr
from transformers import AutoTokenizer, VisionEncoderDecoderModel, ViTImageProcessor, pipeline
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 predict(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=predict(image)
    # Generate an image from the caption
    generated_image = diffusion_model(caption)["sample"][0]

    return caption, generated_images[0]
    
# Set up Gradio interface
input = gr.Image(label="Upload any Image", type='pil')
outputs = [gr.Textbox(label="Caption"), gr.Image(label="Generated Image")]
#examples = ['example1.jpeg']

title = "Image Captioning and Generation"
interface = gr.Interface(
    fn=predict,
    inputs=input,
    outputs=outputs,
    #examples=examples,
    title=title,
)
interface.launch(share=True)