hidhann commited on
Commit
7d8974a
·
verified ·
1 Parent(s): 3c07dcb

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +48 -0
app.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ from PIL import Image
3
+ import torch
4
+ from diffusers import ControlNetModel, StableDiffusionXLControlNetPipeline
5
+ import io
6
+
7
+ @st.cache_resource
8
+ def load_model():
9
+ controlnet = ControlNetModel.from_pretrained("monster-labs/control_v1p_sdxl_qrcode_monster")
10
+ pipeline = StableDiffusionXLControlNetPipeline.from_pretrained(
11
+ "stabilityai/stable-diffusion-xl-base-1.0",
12
+ controlnet=controlnet,
13
+ torch_dtype=torch.float16
14
+ )
15
+ return pipeline.to("cuda" if torch.cuda.is_available() else "cpu")
16
+
17
+ pipeline = load_model()
18
+
19
+ st.title("QR Code Image Generator")
20
+
21
+ uploaded_file = st.file_uploader("Choose a QR code image", type=["png", "jpg", "jpeg"])
22
+ prompt = st.text_input("Enter your prompt")
23
+ num_inference_steps = st.slider("Number of inference steps", 1, 100, 30)
24
+ guidance_scale = st.slider("Guidance scale", 1.0, 20.0, 7.5)
25
+
26
+ if uploaded_file is not None and prompt:
27
+ qr_image = Image.open(uploaded_file).convert("RGB")
28
+
29
+ if st.button("Generate Image"):
30
+ with st.spinner("Generating image..."):
31
+ image = pipeline(
32
+ prompt,
33
+ image=qr_image,
34
+ num_inference_steps=num_inference_steps,
35
+ guidance_scale=guidance_scale
36
+ ).images[0]
37
+
38
+ st.image(image, caption="Generated Image")
39
+
40
+ # Option to download the generated image
41
+ buf = io.BytesIO()
42
+ image.save(buf, format="PNG")
43
+ btn = st.download_button(
44
+ label="Download image",
45
+ data=buf.getvalue(),
46
+ file_name="generated_image.png",
47
+ mime="image/png"
48
+ )