app.py

import gradio as gr
from transformers import pipeline
from diffusers import StableDiffusionPipeline
from diffusers import DiffusionPipeline
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
model_id = "CompVis/stable-diffusion-v1-4"
torch_dtype = torch.float32

if torch.cuda.is_available():
    torch_dtype = torch.bfloat16

def generate_description(image):
    model = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
    return model(image)[0]['generated_text']

def generate_image_by_description(description):
    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch_dtype)
    pipe = pipe.to(device)
    pipe.enable_attention_slicing()

    prompt = (
        f"Generate a image of a pigeon for my profile avatar. "
        f"The description of the pigeon is: {description}. "
        )
    image = pipe(prompt).images[0]  
    return image


with gr.Blocks() as demo:
    with gr.Row():
        with gr.Column(scale=2, min_width=300):
            selected_image = gr.Image(type="filepath", label="Upload an Image of the Pigeon",height=300)
            generate_button = gr.Button("Generate Avatar", variant="primary")
        with gr.Column(scale=2, min_width=300):
            generated_image = gr.Image(type="numpy", label="Generated Avatar", height=300)
            def process_and_generate(image):
                description = generate_description(image)
                return generate_image_by_description(description)

            generate_button.click(process_and_generate, inputs=selected_image, outputs=generated_image)
demo.launch()