app.py

import streamlit as st
from datetime import datetime
import os
from io import BytesIO
from huggingface_hub import InferenceClient

# Initialize InferenceClient with nscale provider
client = InferenceClient(
    provider="nscale",
    api_key=os.environ.get("HF_TOKEN", ""),
)

# Page configuration - using centered layout to fix shaking in wide mode
st.set_page_config(
    page_title="Cloud Image Generator",
    page_icon="🎨",
    layout="centered"
)

# Custom CSS for better styling
st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        color: #FF4B4B;
        text-align: center;
        margin-bottom: 2rem;
    }
    .stButton>button {
        width: 100%;
        background-color: #FF4B4B;
        color: white;
        font-size: 1.2rem;
        padding: 0.5rem 1rem;
        border-radius: 10px;
    }
    .stButton>button:hover {
        background-color: #FF6B6B;
    }
    /* Prevent layout shift */
    .block-container {
        padding-top: 2rem;
        padding-bottom: 2rem;
    }
</style>
""", unsafe_allow_html=True)

# Initialize session state for generated image
if 'generated_image' not in st.session_state:
    st.session_state.generated_image = None
    st.session_state.generated_filename = None

# Header
st.markdown('<h1 class="main-header">🎨 Cloud Image Generator</h1>', unsafe_allow_html=True)
st.markdown("---")

# Main content area
col1, col2 = st.columns([1, 1])

with col1:
    st.header("📝 Prompt")
    
    # Default prompt
    default_prompt = """Develop a 6-frame storyboard for an espresso machine guide following the provided framework, but with a 'modern lifestyle' tone. Imagine this is for a trendy home-appliance brand—sleek, upbeat, and exciting. The visual descriptions should highlight the aesthetic of the process (the golden crema, the steam, the polished chrome), while the narration should emphasize how easy and rewarding it is to craft a cafe-quality drink at home. Strictly adhere to the 6-step structure and output requirements."""
    
    prompt = st.text_area(
        "Enter your image prompt",
        value=default_prompt,
        height=300,
        help="Describe the image you want to generate"
    )
    
    # Generate button
    generate_button = st.button("🎨 Generate Image", type="primary")

with col2:
    st.header("🖼️ Generated Image")
    
    # Placeholder for the image
    image_placeholder = st.empty()
    
    # Display previously generated image if exists
    if st.session_state.generated_image is not None:
        image_placeholder.image(
            st.session_state.generated_image,
            caption=st.session_state.generated_filename,
            use_container_width=True
        )
        
        # Download button
        st.download_button(
            label="📥 Download Image",
            data=st.session_state.generated_image,
            file_name=st.session_state.generated_filename,
            mime="image/png",
            type="primary"
        )
    else:
        image_placeholder.info("👆 Your generated image will appear here")

# Handle image generation
if generate_button:
    if not prompt.strip():
        st.error("Please enter a prompt!")
    else:
        # Create a progress bar
        progress_bar = st.progress(0)
        status_text = st.empty()
        
        try:
            status_text.text("🔄 Connecting to API...")
            progress_bar.progress(20)
            
            # Check if HF_TOKEN is configured
            if not os.environ.get("HF_TOKEN"):
                raise ValueError("HF_TOKEN not configured. Please set the HF_TOKEN secret in Space settings.")
            
            status_text.text("🎨 Generating image...")
            progress_bar.progress(50)
            
            # Use InferenceClient for image generation
            image = client.text_to_image(
                prompt,
                model="stabilityai/stable-diffusion-xl-base-1.0",
            )
            
            progress_bar.progress(80)
            status_text.text("💾 Processing image...")
            
            # Convert PIL Image to bytes
            img_byte_arr = BytesIO()
            image.save(img_byte_arr, format='PNG')
            img_byte_arr.seek(0)
            image_bytes = img_byte_arr.getvalue()
            
            # Generate filename
            timestamp = datetime.now().timestamp()
            filename = f"generated-image-{int(timestamp)}.png"
            
            # Store image in session state
            st.session_state.generated_image = image_bytes
            st.session_state.generated_filename = filename
            
            progress_bar.progress(100)
            status_text.text("✅ Image generated successfully!")
            
            # Display the image
            image_placeholder.image(
                image_bytes,
                caption=filename,
                use_container_width=True
            )
            
            st.success("🎉 Image generated successfully!")
            
            # Download button
            st.download_button(
                label="📥 Download Image",
                data=image_bytes,
                file_name=filename,
                mime="image/png",
                type="primary"
            )
            
        except Exception as e:
            progress_bar.empty()
            status_text.empty()
            st.error(f"❌ Error: {str(e)}")

# Footer
st.markdown("---")
st.markdown(
    "<p style='text-align: center; color: gray;'>Made with ❤️ using Streamlit</p>",
    unsafe_allow_html=True
)