Create app.py

app.py

@@ -0,0 +1,96 @@
+import streamlit as st
+from diffusers import DiffusionPipeline
+import torch
+
+# Set page config
+st.set_page_config(
+    page_title="Portrait Generator",
+    page_icon="🖼️",
+    layout="centered"
+)
+
+# App title and description
+st.title("AI Portrait Generator")
+st.markdown("Generate beautiful portraits using the AWPortraitCN2 model")
+
+# Model parameters
+with st.sidebar:
+    st.header("Generation Settings")
+    steps = st.slider("Inference Steps", min_value=20, max_value=100, value=40)
+    guidance_scale = st.slider("Guidance Scale", min_value=1.0, max_value=15.0, value=7.5, step=0.5)
+    negative_prompt = st.text_area(
+        "Negative Prompt",
+        value="lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, watermark, signature, out of frame"
+    )
+    seed = st.number_input("Random Seed (leave at -1 for random)", min_value=-1, value=-1)
+    
+# Main prompt input
+prompt = st.text_area(
+    "Describe the portrait you want to generate",
+    value="Masterpiece portrait of a beautiful young woman with flowing hair, detailed face, photorealistic, 8k, professional photography"
+)
+
+# Generate button
+if st.button("Generate Portrait", type="primary"):
+    with st.spinner("Loading model and generating portrait..."):
+        try:
+            # 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)
+            
+            # Set seed if specified
+            generator = None
+            if seed != -1:
+                generator = torch.Generator(device).manual_seed(seed)
+            
+            # Generate the image
+            image = pipeline(
+                prompt=prompt,
+                negative_prompt=negative_prompt,
+                num_inference_steps=steps,
+                guidance_scale=guidance_scale,
+                generator=generator
+            ).images[0]
+            
+            # Display the generated image
+            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()
+            
+            st.download_button(
+                label="Download Portrait",
+                data=byte_im,
+                file_name="generated_portrait.png",
+                mime="image/png"
+            )
+            
+        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.")
+
+# 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`
+""")