Spaces:
Running
Running
| import streamlit as st | |
| import torch | |
| from PIL import Image | |
| import io | |
| # 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" | |
| ) | |
| # 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", | |
| 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 | |
| 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` | |
| """) |