nimra-mughal commited on
Commit
7168107
·
verified ·
1 Parent(s): 4b851e7

Update src/streamlit_app.py

Browse files
Files changed (1) hide show
  1. src/streamlit_app.py +7 -3
src/streamlit_app.py CHANGED
@@ -1,6 +1,12 @@
1
  import streamlit as st
2
  from diffusers import StableDiffusionPipeline
3
  import torch
 
 
 
 
 
 
4
 
5
  st.set_page_config(page_title="AI Design Generator", layout="centered")
6
 
@@ -10,7 +16,7 @@ def load_model():
10
  "runwayml/stable-diffusion-v1-5",
11
  torch_dtype=torch.float32
12
  )
13
- model.to("cpu") # Space is on CPU unless you enable GPU manually
14
  return model
15
 
16
  pipe = load_model()
@@ -22,5 +28,3 @@ if st.button("Generate Image"):
22
  with st.spinner("Generating..."):
23
  image = pipe(prompt).images[0]
24
  st.image(image, caption="Generated Design", use_column_width=True)
25
-
26
-
 
1
  import streamlit as st
2
  from diffusers import StableDiffusionPipeline
3
  import torch
4
+ import os
5
+
6
+ # Set cache to writable directory
7
+ os.environ["HF_HOME"] = "/tmp"
8
+ os.environ["TRANSFORMERS_CACHE"] = "/tmp"
9
+ os.environ["HF_HUB_CACHE"] = "/tmp"
10
 
11
  st.set_page_config(page_title="AI Design Generator", layout="centered")
12
 
 
16
  "runwayml/stable-diffusion-v1-5",
17
  torch_dtype=torch.float32
18
  )
19
+ model.to("cpu") # CPU unless GPU enabled in Space settings
20
  return model
21
 
22
  pipe = load_model()
 
28
  with st.spinner("Generating..."):
29
  image = pipe(prompt).images[0]
30
  st.image(image, caption="Generated Design", use_column_width=True)