Spaces:
Build error
Build error
File size: 2,902 Bytes
d852a69 f2eb806 16ffc06 f2eb806 d852a69 f2eb806 16ffc06 f2eb806 16ffc06 f2eb806 d852a69 f2eb806 d852a69 f2eb806 d852a69 f2eb806 d852a69 f2eb806 d852a69 | 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 | import streamlit as st
from diffusers import StableDiffusionPipeline
from PIL import Image
import torch
import random
import os
# Streamlit UI configuration - This must be the first Streamlit command
st.set_page_config(page_title="AI Image Variation Generator", layout="wide")
# Set device (CPU/GPU) and appropriate dtype
device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32
# Load the Stable Diffusion XL model with mixed precision if on GPU
sd_pipe = StableDiffusionPipeline.from_pretrained(
"stabilityai/stable-diffusion-xl-base-1.0",
torch_dtype=torch_dtype
)
sd_pipe.to(device)
# Enable memory-efficient attention if available; fallback to attention slicing
try:
sd_pipe.enable_xformers_memory_efficient_attention()
except Exception as e:
st.warning("xFormers memory efficient attention not available. Falling back to attention slicing.")
sd_pipe.enable_attention_slicing()
st.title("🎨 AI Image Variation Generator")
st.write("Upload an image and get AI-generated variations based on your style instructions!")
uploaded_file = st.file_uploader("Upload an Image", type=["png", "jpg", "jpeg"])
prompt = st.text_input("Enter style modification instructions (optional):")
if uploaded_file:
image = Image.open(uploaded_file).convert("RGB")
st.image(image, caption="Uploaded Image", use_column_width=True)
if st.button("Generate Variations"):
with st.spinner("Generating variations..."):
variations = []
# Use torch.no_grad() to reduce memory overhead during inference
with torch.no_grad():
for i in range(4): # Generate 4 variations
seed = random.randint(0, 10000)
generator = torch.manual_seed(seed)
# Generate image using the pipeline with specified prompt and steps
img = sd_pipe(prompt=prompt, num_inference_steps=30, generator=generator).images[0]
filename = f"variation_{i+1}.png"
img.save(filename)
variations.append((img, filename))
# Clear GPU cache to free up memory between iterations
if device == "cuda":
torch.cuda.empty_cache()
# Display the generated variations in two columns
col1, col2 = st.columns(2)
for idx, (img, filename) in enumerate(variations):
with (col1 if idx % 2 == 0 else col2):
st.image(img, caption=f"Variation {idx+1}", use_column_width=True)
# Provide a download button for each variation
with open(filename, "rb") as file:
st.download_button("Download", file.read(), file_name=filename)
if st.button("Regenerate"):
st.experimental_rerun()
|