File size: 5,035 Bytes
f5349b6
 
 
 
 
 
 
 
 
 
 
 
 
0da87cb
 
f5349b6
 
 
 
 
 
 
 
 
 
 
 
 
92b7960
f5349b6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fefe39a
 
f5349b6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3db6d75
f5349b6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fefe39a
f5349b6
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
fefe39a
 
f5349b6
 
 
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
import streamlit as st
import requests
import io
import os
from PIL import Image
from datetime import datetime

# Set up Hugging Face API key
HUGGINGFACE_API_KEY = os.getenv('HUGGINGFACE_API_KEY')

# Dictionary of available models and their API endpoints
MODELS = {
    "Stable Diffusion XL": "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0",
    "Flux 1": "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev",
    "Flux Midjourney": "https://api-inference.huggingface.co/models/strangerzonehf/Flux-Midjourney-Mix2-LoRA",
    "Playground v2": "https://api-inference.huggingface.co/models/playgroundai/playground-v2-1024px-aesthetic"
}

def generate_image(prompt, model_url):
    """Generate image using Hugging Face API"""
    headers = {
        "Authorization": f"Bearer {HUGGINGFACE_API_KEY}",
        "Content-Type": "application/json"
    }
    
    data = {
        "inputs": prompt,
        "parameters": {
            "num_inference_steps": 20,
            "guidance_scale": 7.5
        }
    }
    
    response = requests.post(
        model_url,
        headers=headers,
        json=data
    )
    
    if response.status_code == 200:
        return response.content
    else:
        st.error(f"Error: {response.text}")
        return None

def save_image(image_data, model_name):
    """Save the generated image"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"generated_{model_name}_{timestamp}.png"
    
    image = Image.open(io.BytesIO(image_data))
    image.save(filename)
    return filename

# Streamlit UI
st.title("Open Source Text-to-Image Generation")
st.markdown("""

This app uses state-of-the-art open source text-to-image models via Hugging Face's API.

- **SDXL**: Latest Stable Diffusion model with exceptional quality

- **Flux 1**: Optimized for speed and quality

- **Flux Midjourney**: Compact but powerful model

- **Playground v2**: Specialized in aesthetic generations

""")

# API key input
if 'api_key_set' not in st.session_state:
    st.session_state.api_key_set = False

if not st.session_state.api_key_set:
    with st.form("api_key_form"):
        huggingface_key = st.text_input("Enter Hugging Face API Key:", type="password", 
                                       help="Get your free API key from huggingface.co")
        submit = st.form_submit_button("Save API Key")
        
        if submit:
            os.environ['HUGGINGFACE_API_KEY'] = huggingface_key
            st.session_state.api_key_set = True
            st.rerun()

if st.session_state.api_key_set:
    # Model selection
    model = st.selectbox(
        "Choose Model:",
        list(MODELS.keys())
    )
    
    # Advanced options
    with st.expander("Advanced Options"):
        seed = st.number_input("Random Seed (optional)", min_value=0, max_value=999999999, value=0)
        if seed > 0:
            st.info("Using the same seed will generate similar images for the same prompt")

    # Prompt input
    prompt = st.text_area(
        "Enter your prompt:",
        height=100,
        help="Describe what you want to see in the image. Be specific!"
    )
    
    # Generate image button

    if st.button("Generate Image"):
        if prompt:
            with st.spinner(f"Generating image using {model}..."):
                image_data = generate_image(prompt, MODELS[model])
                
                if image_data:
                    # Save the image
                    filename = save_image(image_data, model.lower().replace(" ", "_"))
                    
                    # Display the image
                    st.image(Image.open(io.BytesIO(image_data)))
                    
                    # Display prompt used
                    st.text_area("Prompt used:", prompt, height=100, disabled=True)
                    
                    # Download button
                    with open(filename, "rb") as file:
                        btn = st.download_button(
                            label="Download Image",
                            data=file,
                            file_name=filename,
                            mime="image/png"
                        )
        else:
            st.warning("Please enter a prompt.")

    # Display model information
    st.markdown("### About the Selected Model")
    model_info = {
        "Stable Diffusion XL": "The latest version of Stable Diffusion, known for its high-quality outputs and improved understanding of prompts.",
        "Flux 1": "A highly optimized model that balances speed and quality, great for quick iterations.",
        "Flux Midjourney": "A compact 1B parameter model that produces impressive results with faster inference times.",
        "Playground v2": "Specialized in creating highly aesthetic images with strong composition and artistic quality."
    }
    st.info(model_info[model])