Spaces:
Sleeping
Sleeping
File size: 1,884 Bytes
6620294 2709dad 6620294 2709dad 6620294 2709dad 33799cd 6620294 | 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 | 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()
|