Spaces:
Build error
Build error
| 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() | |