Spaces:
Running
Running
File size: 4,735 Bytes
b9dd8bb f5940c8 5655d9c b9dd8bb 5655d9c b9dd8bb 5655d9c b9dd8bb 5655d9c f5940c8 5655d9c b9dd8bb 5655d9c b9dd8bb 5655d9c b9dd8bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 |
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(
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")
# 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")
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"
)
# 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:
# Load the model
pipeline = load_model()
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
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
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("""
### 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.
""") |