Spaces:
Running
Running
import streamlit as st | |
from PIL import Image | |
import torch | |
from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline | |
import io | |
def load_model(): | |
controlnet = ControlNetModel.from_pretrained("monster-labs/control_v1p_sdxl_qrcode_monster") | |
pipeline = StableDiffusionXLControlNetPipeline.from_pretrained( | |
"stabilityai/stable-diffusion-xl-base-1.0", | |
controlnet=controlnet, | |
torch_dtype=torch.float16 | |
) | |
return pipeline.to("cuda" if torch.cuda.is_available() else "cpu") | |
pipeline = load_model() | |
st.title("QR Code Image Generator") | |
uploaded_file = st.file_uploader("Choose a QR code image", type=["png", "jpg", "jpeg"]) | |
prompt = st.text_input("Enter your prompt") | |
num_inference_steps = st.slider("Number of inference steps", 1, 100, 30) | |
guidance_scale = st.slider("Guidance scale", 1.0, 20.0, 7.5) | |
if uploaded_file is not None and prompt: | |
qr_image = Image.open(uploaded_file).convert("RGB") | |
if st.button("Generate Image"): | |
with st.spinner("Generating image..."): | |
image = pipeline( | |
prompt, | |
image=qr_image, | |
num_inference_steps=num_inference_steps, | |
guidance_scale=guidance_scale | |
).images[0] | |
st.image(image, caption="Generated Image") | |
# Option to download the generated image | |
buf = io.BytesIO() | |
image.save(buf, format="PNG") | |
btn = st.download_button( | |
label="Download image", | |
data=buf.getvalue(), | |
file_name="generated_image.png", | |
mime="image/png" | |
) |