Updated app.py

app.py

@@ -2,6 +2,9 @@ import streamlit as st
 import torch
 from PIL import Image
 import io
+import os
+import subprocess
+import sys
 
 # Set page config
 st.set_page_config(
@@ -14,6 +17,36 @@ st.set_page_config(
 st.title("AI Portrait Generator")
 st.markdown("Generate beautiful portraits using the AWPortraitCN2 model")
 
+# Check and install compatible versions if needed
+@st.cache_resource
+def install_dependencies():
+    try:
+        # Try to import diffusers to see if it works
+        import diffusers
+        return True
+    except ImportError:
+        st.warning("Installing required packages. This may take a few minutes...")
+        # Install specific versions known to work together
+        subprocess.check_call([
+            sys.executable, "-m", "pip", "install", 
+            "huggingface-hub==0.16.4", 
+            "diffusers==0.20.0",
+            "transformers==4.32.0",
+            "accelerate==0.21.0"
+        ])
+        return True
+    except Exception as e:
+        st.error(f"Failed to install dependencies: {e}")
+        return False
+
+# Try to install compatible dependencies
+dependencies_installed = install_dependencies()
+
+# If dependencies installation failed, show message and exit
+if not dependencies_installed:
+    st.error("Could not set up the required environment. Please check the logs.")
+    st.stop()
+
 # Model parameters
 with st.sidebar:
     st.header("Generation Settings")
@@ -31,27 +64,42 @@ prompt = st.text_area(
     value="Masterpiece portrait of a beautiful young woman with flowing hair, detailed face, photorealistic, 8k, professional photography"
 )
 
-# Load the model when the user clicks the button
+# Function to load model with proper dependencies
+@st.cache_resource
+def load_model():
+    try:
+        from diffusers import StableDiffusionPipeline
+        
+        pipeline = StableDiffusionPipeline.from_pretrained(
+            "Shakker-Labs/AWPortraitCN2",
+            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            use_safetensors=True
+        )
+        
+        # Move to GPU if available
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        pipeline = pipeline.to(device)
+        
+        return pipeline
+    except Exception as e:
+        st.error(f"Error loading model: {e}")
+        return None
+
+# Generate 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
+            # Load the model
+            pipeline = load_model()
             
-            # Set up the model pipeline
-            pipeline = DiffusionPipeline.from_pretrained(
-                "Shakker-Labs/AWPortraitCN2",
-                torch_dtype=torch.float16,
-                use_safetensors=True
-            )
-            
-            # Move to GPU if available
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-            pipeline = pipeline.to(device)
+            if pipeline is None:
+                st.error("Failed to load the model. Please check the logs.")
+                st.stop()
             
             # Set seed if specified
             generator = None
             if seed != -1:
+                device = "cuda" if torch.cuda.is_available() else "cpu"
                 generator = torch.Generator(device).manual_seed(seed)
             
             # Generate the image
@@ -85,13 +133,7 @@ if st.button("Generate Portrait", type="primary"):
 # Add requirements info at the bottom
 st.markdown("---")
 st.markdown("""
-### Requirements
-To run this app, you need:
-- diffusers
-- transformers
-- accelerate
-- torch
-- streamlit
-
-Install with: `pip install diffusers transformers accelerate torch streamlit`
+### About This App
+This app uses the AWPortraitCN2 model to generate AI portraits based on your text prompts.
+Adjust the settings in the sidebar to customize your generation.
 """)