Updated app.py

app.py

@@ -1,10 +1,8 @@
 import streamlit as st
 import torch
-from PIL import Image
 import io
 import os
-import subprocess
-import sys
+from PIL import Image
 
 # Set page config
 st.set_page_config(
@@ -17,36 +15,6 @@ 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")
@@ -64,16 +32,19 @@ prompt = st.text_area(
     value="Masterpiece portrait of a beautiful young woman with flowing hair, detailed face, photorealistic, 8k, professional photography"
 )
 
-# Function to load model with proper dependencies
+# Function to load model using modern API
 @st.cache_resource
 def load_model():
     try:
-        from diffusers import StableDiffusionPipeline
+        # Import these inside the function to handle errors gracefully
+        from diffusers import AutoPipelineForText2Image
         
-        pipeline = StableDiffusionPipeline.from_pretrained(
+        # Use AutoPipeline which is more compatible with newer versions
+        pipeline = AutoPipelineForText2Image.from_pretrained(
             "Shakker-Labs/AWPortraitCN2",
             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            use_safetensors=True
+            use_safetensors=True,
+            variant="fp16" if torch.cuda.is_available() else None
         )
         
         # Move to GPU if available
@@ -82,8 +53,28 @@ def load_model():
         
         return pipeline
     except Exception as e:
-        st.error(f"Error loading model: {e}")
-        return None
+        st.error(f"Error loading model: {str(e)}")
+        st.info("Debug info: Using modern API with AutoPipelineForText2Image")
+        
+        # Fallback to traditional StableDiffusionPipeline if needed
+        try:
+            st.info("Trying alternative method...")
+            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 e2:
+            st.error(f"Alternative method also failed: {str(e2)}")
+            return None
 
 # Generate button
 if st.button("Generate Portrait", type="primary"):
@@ -93,7 +84,7 @@ if st.button("Generate Portrait", type="primary"):
             pipeline = load_model()
             
             if pipeline is None:
-                st.error("Failed to load the model. Please check the logs.")
+                st.error("Failed to load the model. Check the logs for details.")
                 st.stop()
             
             # Set seed if specified
@@ -127,13 +118,17 @@ if st.button("Generate Portrait", type="primary"):
             )
             
         except Exception as e:
-            st.error(f"An error occurred: {e}")
-            st.info("Make sure you have enough GPU memory and the required dependencies installed.")
+            st.error(f"An error occurred during generation: {str(e)}")
+            st.info("Make sure you have enough GPU memory (T4 or better recommended).")
 
-# Add requirements info at the bottom
-st.markdown("---")
-st.markdown("""
-### 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.
-""")
+# Add hardware info at the bottom
+if torch.cuda.is_available():
+    st.markdown("---")
+    st.markdown(f"""
+    ### Hardware Info
+    - Running on: {torch.cuda.get_device_name(0)}
+    - Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.2f} GB
+    """)
+else:
+    st.markdown("---")
+    st.markdown("⚠️ Running on CPU. For better performance, use a GPU runtime.")