File size: 3,603 Bytes
2a9628e |
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 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import streamlit as st
from PIL import Image
import numpy as np
import io
import torch
from diffusers import DiffusionPipeline
from diffusers import StableDiffusionPipeline
# Function to generate images based on user input
def generate_images(text, size, num_images, quality):
# Placeholder function, replace with your actual implementation
generated_images = []
for _ in range(num_images):
# Use your model to generate an image based on the input text
# Example:
# image = model.generate_image(text, size, quality)
# generated_images.append(image)
pass
return generated_images
# Function to convert PIL image to bytes
def pil_to_bytes(image):
img_byte_array = io.BytesIO()
image.save(img_byte_array, format='PNG')
return img_byte_array.getvalue()
# Streamlit interface
st.title('Text-to-Image AI-based Chatbot')
# Text input box
user_input = st.text_input("Enter your text here:")
# Select models
model_opts = ["amused/amused-256", "amused/amused-512", "runwayml/stable-diffusion-v1-5"]
selected_model = st.selectbox("Select model:", model_opts)
if selected_model == 'amused/amused-256':
print("Model selected: 256")
pipe = DiffusionPipeline.from_pretrained('amused/amused-256', variant="fp16", torch_dtype=torch.float16)
elif selected_model == 'amused/amused-512':
print("Model selected: 512")
pipe = DiffusionPipeline.from_pretrained('amused/amused-512', variant="fp16", torch_dtype=torch.float16)
else:
print("Model selected: diffusion")
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", variant="fp16", torch_dtype=torch.float16)
# # Image size input
# image_size_options = ["1080x1080", "1280x720", "1920x1080"]
# selected_size = st.selectbox("Select image size:", image_size_options)
# Convert selected size to width and height
# width, height = map(int, selected_size.split('x'))
# Number of images input
num_images = st.number_input("Number of images to generate:", min_value=1, max_value=10, value=3, step=1)
# Image quality input
image_quality = st.slider("Select image quality:", min_value=1, max_value=100, value=50, step=1)
# Generate button
if st.button("Generate Images"):
if user_input:
# Example image
# example_image = Image.open("example_img.jpg")
# st.image(example_image, caption='Example Image', use_column_width=True)
# img_bytes = pil_to_bytes(example_image)
# st.download_button(label=f"Download Image ", data=img_bytes, file_name=f"generated_image_example.png", mime='image/png')
pipe = pipe.to('cuda')
with st.spinner('Generating...'):
image = pipe(
user_input,
num_images_per_prompt=num_images,
num_inference_steps=image_quality,
generator=torch.Generator('cuda').manual_seed(8)
).images
st.image(image , use_column_width=True)
# img_bytes = pil_to_bytes(image)
# st.download_button(label=f"Download Image ", data=img_bytes, file_name=f"generated_image_example.png", mime='image/png')
# # Generate images based on user input
# generated_images = generate_images(user_input, (width, height), num_images, image_quality)
# if generated_images:
# for i, img in enumerate(generated_images):
# st.image(img, caption=f"Generated Image {i+1}", use_column_width=True)
# # Download link for each image
# img_bytes = pil_to_bytes(img)
# st.download_butt |