Updated app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
-from diffusers import DiffusionPipeline
 import torch
+from PIL import Image
+import io
 
 # Set page config
 st.set_page_config(
@@ -30,10 +31,13 @@ prompt = st.text_area(
     value="Masterpiece portrait of a beautiful young woman with flowing hair, detailed face, photorealistic, 8k, professional photography"
 )
 
-# Generate button
+# Load the model when the user clicks the button
 if st.button("Generate Portrait", type="primary"):
     with st.spinner("Loading model and generating portrait..."):
         try:
+            # Import here to avoid the cached_download issue until the model is actually needed
+            from diffusers import DiffusionPipeline
+            
             # Set up the model pipeline
             pipeline = DiffusionPipeline.from_pretrained(
                 "Shakker-Labs/AWPortraitCN2",
@@ -63,9 +67,6 @@ if st.button("Generate Portrait", type="primary"):
             st.image(image, caption="Generated Portrait", use_column_width=True)
             
             # Option to download
-            # Convert the PIL image to bytes
-            import io
-            from PIL import Image
             buf = io.BytesIO()
             image.save(buf, format="PNG")
             byte_im = buf.getvalue()