Spaces:
Sleeping
Sleeping
File size: 2,145 Bytes
3cc83d8 a1de4cd 3cc83d8 2e2572c 3cc83d8 |
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 |
import streamlit as st
from PIL import Image
from io import BytesIO
import base64
from diffusers import StableDiffusionPipeline
import torch
# Initialize the Stable Diffusion model
model_id = "stabilityai/stable-diffusion-3-medium"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
pipe.to("cpu")
def generate_image(prompt, negative_prompt=None, temperature=1.0, steps=50, image_size=(512, 512)):
# Generate an image using the Stable Diffusion pipeline
image = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=steps, guidance_scale=temperature).images[0]
# Resize image
image = image.resize(image_size)
# Convert image to base64
buffered = BytesIO()
image.save(buffered, format="PNG")
img_str = base64.b64encode(buffered.getvalue()).decode('utf-8')
return img_str
def main():
st.title("Stable Diffusion Image Generation API")
st.write("Generate images using Stable Diffusion and get them in base64 format.")
# Get parameters from URL
query_params = st.experimental_get_query_params()
prompt = query_params.get("prompt", [""])[0]
negative_prompt = query_params.get("negative_prompt", [None])[0]
temperature = float(query_params.get("temperature", [1.0])[0])
steps = int(query_params.get("steps", [50])[0])
image_size = tuple(map(int, query_params.get("image_size", ["512,512"])[0].split(",")))
if prompt:
st.write("Generating image with parameters:")
st.write(f"Prompt: {prompt}")
st.write(f"Negative Prompt: {negative_prompt}")
st.write(f"Temperature: {temperature}")
st.write(f"Steps: {steps}")
st.write(f"Image Size: {image_size}")
# Generate the image
img_base64 = generate_image(prompt, negative_prompt, temperature, steps, image_size)
# Display the image
st.image(f"data:image/png;base64,{img_base64}", caption="Generated Image")
# Provide the base64 image string
st.text_area("Base64 Image String", value=img_base64, height=200)
if __name__ == "__main__":
main() |