text2image / app.py
My-AI-Projects's picture
Update app.py
3b42df3 verified
raw
history blame
868 Bytes
import gradio as gr
from diffusers import DiffusionPipeline
import torch
# Load the Kwai-Kolors diffusion model
pipeline = DiffusionPipeline.from_pretrained("Kwai-Kolors/Kolors").to("cuda" if torch.cuda.is_available() else "cpu")
# Function to generate image from prompt
def generate_image(prompt):
# Use the pipeline to generate an image from the text prompt
image = pipeline(prompt).images[0]
return image
# Create Gradio interface
iface = gr.Interface(
fn=generate_image, # Function that generates the image
inputs=gr.Textbox(lines=2, placeholder="Enter your prompt"), # Textbox input for the prompt
outputs="image", # Output is an image
title="Kwai-Kolors Image Generator",
description="Generate images from text prompts using the Kwai-Kolors diffusion model."
)
# Launch the app
iface.launch()