import gradio as gr
from transformers import pipeline
from diffusers import StableDiffusionPipeline
import clip
import torch
from PIL import Image

device = "cuda" if torch.cuda.is_available() else "cpu"

# Load models from your Hugging Face Hub repos
text_generator = pipeline("text-generation", model="roshanVarghese/my-gpt2-model")
image_pipe = StableDiffusionPipeline.from_pretrained("roshanVarghese/my-stable-diffusion").to(device)
clip_model, clip_preprocess = clip.load("ViT-B/32", device=device)

def generate_story(prompt):
    # 'max_length' works better for pipeline than 'max_new_tokens'
    return text_generator(prompt, max_length=200, do_sample=True)[0]['generated_text']

def generate_image(story):
    image = image_pipe(story).images[0]
    image.save("generated_image.png")
    return image

def evaluate_similarity(story, img_path="generated_image.png"):
    image = clip_preprocess(Image.open(img_path)).unsqueeze(0).to(device)
    text = clip.tokenize([story], truncate=True).to(device)
    with torch.no_grad():
        image_features = clip_model.encode_image(image)
        text_features = clip_model.encode_text(text)
        similarity = torch.nn.functional.cosine_similarity(image_features, text_features).item()
    return similarity

def gradio_pipeline(prompt):
    story = generate_story(prompt)
    image = generate_image(story)
    score = evaluate_similarity(story)
    return story, image, f"Similarity Score: {score:.2f}"

iface = gr.Interface(
    fn=gradio_pipeline,
    inputs=gr.Textbox(label="Enter a Story Prompt"),
    outputs=[gr.Textbox(label="Generated Story"), gr.Image(label="Generated Image"), gr.Textbox(label="Image-Story Similarity")],
    title="Story-to-Image AI Pipeline",
    description="Enter a prompt. The AI will generate a story, create an image, and evaluate similarity."
)

if __name__ == "__main__":
    iface.launch()