app.py

import os
import sys

# Set critical environment variables first
os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
os.environ["WATCHDOG_OPTIONAL"] = "1"
os.environ["PYTORCH_JIT"] = "0"

# Import third party modules
import streamlit as st
import numpy as np
import random
from PIL import Image
import io
import time

# Set up imports for huggingface_hub
# Import what we can, but handle potential import errors
try:
    from huggingface_hub import HfApi, HfFolder, login
except ImportError as e:
    st.error(f"Error importing from huggingface_hub: {e}")

# Configure Hugging Face cache and environment
os.environ["HF_HOME"] = os.path.join(os.getcwd(), ".cache/huggingface")

# Import PyTorch after environment setup
import torch
from diffusers import FluxFillPipeline

import warnings
warnings.filterwarnings("ignore", message=".*add_prefix_space.*")

# Constants
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

# Setting page config
st.set_page_config(
    page_title="FLUX.1 Fill [dev]",
    layout="wide"
)

# Title and description
st.markdown("""
# FLUX.1 Fill [dev]
12B param rectified flow transformer structural conditioning tuned, guidance-distilled from [FLUX.1 [pro]](https://blackforestlabs.ai/)  
[[non-commercial license](https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/LICENSE.md)] [[blog](https://blackforestlabs.ai/announcing-black-forest-labs/)] [[model](https://huggingface.co/black-forest-labs/FLUX.1-dev)]
""")

# Add simple instructions
st.sidebar.markdown("""
## Important Setup Information

This app uses the FLUX.1-Fill-dev model which requires special access:

1. Sign up/login at [Hugging Face](https://huggingface.co/)
2. Request access to [FLUX.1-Fill-dev](https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev) by clicking 'Access repository'
3. Wait for approval from model owners

### For Hugging Face Spaces Setup:
1. Go to your Space settings > Secrets
2. Add a new secret with the name `HF_TOKEN` 
3. Set its value to your Hugging Face API token (found in your account settings)
""")

# Try to get a Hugging Face token from environment variables
def get_hf_token():
    # Check common environment variable names for HF tokens
    token_env_vars = [
        'HF_TOKEN', 
        'HUGGINGFACE_TOKEN', 
        'HUGGING_FACE_HUB_TOKEN',
        'HF_API_TOKEN',
        'HUGGINGFACE_API_TOKEN',
        'HUGGINGFACE_HUB_TOKEN'
    ]
    
    for env_var in token_env_vars:
        if env_var in os.environ and os.environ[env_var].strip():
            st.sidebar.success(f"Found token in {env_var}")
            return os.environ[env_var].strip()
    
    # If we're here, no token was found
    st.sidebar.warning("No Hugging Face token found in environment variables")
    
    return None

@st.cache_resource(show_spinner=False)
def load_model():
    """Load the model with a simplified approach using the required token"""
    # Get device
    device = "cuda" if torch.cuda.is_available() else "cpu"
    st.info(f"Using device: {device}")
    
    # Get token
    token = get_hf_token()
    st.info(f"Token available: {'Yes' if token else 'No'}")
    
    try:
        # Add use_fast_tokenizer=True to address the tokenizer warning
        model = FluxFillPipeline.from_pretrained(
            "black-forest-labs/FLUX.1-Fill-dev",
            token=token,
            torch_dtype=torch.bfloat16,
            revision="main"
        )
        st.success("Model loaded successfully!")
        return model.to(device)
    except Exception as e:
        st.error(f"Failed to load model: {e}")
        
        if "401" in str(e) or "access" in str(e).lower() or "denied" in str(e).lower():
            st.error("""
            Access Denied: You need to:
            1. Request access to the model at https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev
            2. Set up your Hugging Face token in Spaces:
               - Go to your Space settings > Secrets
               - Add a new secret with name 'HF_TOKEN'
               - Set its value to your Hugging Face API token
            3. Wait for approval from model owners
            
            Note: You can find your token at https://huggingface.co/settings/tokens
            """)
        st.stop()
        
    except Exception as e:
        st.error(f"Failed to load model after all attempts: {e}")
        
        if "401" in str(e) or "access" in str(e).lower() or "denied" in str(e).lower():
            st.error("""
            Access Denied: You need to:
            1. Request access to the model at https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev
            2. Set up your Hugging Face token in Spaces:
               - Go to your Space settings > Secrets
               - Add a new secret with name 'HF_TOKEN'
               - Set its value to your Hugging Face API token
            3. Wait for approval from model owners
            
            Note: You can find your token at https://huggingface.co/settings/tokens
            """)
        elif "Tried to instantiate class" in str(e):
            st.error("""
            PyTorch class initialization error. Try restarting the app.
            If the error persists, try accessing the app from a different browser.
            """)
        st.stop()

# Initialize model section
with st.spinner("Loading model..."):
    try:
        pipe = load_model()
        st.success("Model loaded successfully!")
    except Exception as e:
        st.error(f"Failed to load model: {str(e)}")
        st.stop()

def calculate_optimal_dimensions(image: Image.Image):
    # Extract the original dimensions
    original_width, original_height = image.size
    
    # Set constants
    MIN_ASPECT_RATIO = 9 / 16
    MAX_ASPECT_RATIO = 16 / 9
    FIXED_DIMENSION = 1024

    # Calculate the aspect ratio of the original image
    original_aspect_ratio = original_width / original_height

    # Determine which dimension to fix
    if original_aspect_ratio > 1:  # Wider than tall
        width = FIXED_DIMENSION
        height = round(FIXED_DIMENSION / original_aspect_ratio)
    else:  # Taller than wide
        height = FIXED_DIMENSION
        width = round(FIXED_DIMENSION * original_aspect_ratio)

    # Ensure dimensions are multiples of 8
    width = (width // 8) * 8
    height = (height // 8) * 8

    # Enforce aspect ratio limits
    calculated_aspect_ratio = width / height
    if calculated_aspect_ratio > MAX_ASPECT_RATIO:
        width = (height * MAX_ASPECT_RATIO // 8) * 8
    elif calculated_aspect_ratio < MIN_ASPECT_RATIO:
        height = (width / MIN_ASPECT_RATIO // 8) * 8

    # Ensure width and height remain above the minimum dimensions
    width = max(width, 576) if width == FIXED_DIMENSION else width
    height = max(height, 576) if height == FIXED_DIMENSION else height

    return width, height

# Create two columns for layout
col1, col2 = st.columns([1, 1])

with col1:
    # Upload image
    uploaded_file = st.file_uploader("Upload an image for inpainting", type=["jpg", "jpeg", "png"])
    
    if uploaded_file is not None:
        # Display the uploaded image
        image = Image.open(uploaded_file).convert("RGB")
        st.image(image, caption="Uploaded Image", use_container_width=True)
        
        # Simple approach to create a mask - select a square area
        st.write("Select an area to inpaint:")
        
        # Get image dimensions
        img_width, img_height = image.size
        
        # Scale for display while maintaining aspect ratio
        display_height = 600
        display_width = int(img_width * (display_height / img_height))
        
        # Create sliders for selecting the area
        col_sliders1, col_sliders2 = st.columns(2)
        
        with col_sliders1:
            x1 = st.slider("Left edge (X1)", 0, img_width, img_width // 4)
            y1 = st.slider("Top edge (Y1)", 0, img_height, img_height // 4)
        
        with col_sliders2:
            x2 = st.slider("Right edge (X2)", x1, img_width, min(x1 + img_width // 2, img_width))
            y2 = st.slider("Bottom edge (Y2)", y1, img_height, min(y1 + img_height // 2, img_height))
        
        # Create a copy of the image to show the mask
        preview_img = image.copy()
        preview_mask = Image.new("L", image.size, 0)
        
        # Draw a white rectangle on the mask
        from PIL import ImageDraw
        draw = ImageDraw.Draw(preview_mask)
        draw.rectangle([(x1, y1), (x2, y2)], fill=255)
        
        # Show the mask on the image
        masked_preview = image.copy()
        # Add semi-transparent white overlay 
        overlay = Image.new("RGBA", image.size, (255, 255, 255, 128))
        masked_preview.paste(overlay, (0, 0), preview_mask)
        
        st.image(masked_preview, caption="Area to inpaint (white overlay)", use_container_width=True)
        
        # Prompt input
        prompt = st.text_input("Enter your prompt")
        
        # Example prompts
        examples = [
            "a tiny astronaut hatching from an egg on the moon",
            "a cat holding a sign that says hello world",
            "an anime illustration of a wiener schnitzel",
        ]
        
        example_prompt = st.selectbox("Or select an example prompt", [""] + examples)
        if example_prompt and not prompt:
            prompt = example_prompt
        
        # Advanced settings with expander
        with st.expander("Advanced Settings"):
            randomize_seed = st.checkbox("Randomize seed", value=True)
            
            if not randomize_seed:
                seed = st.slider("Seed", 0, MAX_SEED, 0)
            else:
                seed = random.randint(0, MAX_SEED)
            
            guidance_scale = st.slider("Guidance Scale", 1.0, 30.0, 3.5, 0.5)
            num_inference_steps = st.slider("Number of inference steps", 1, 50, 28)
        
        # Run button
        run_button = st.button("Generate")

with col2:
    if uploaded_file is not None:
        st.write("Result will appear here")
        
        if run_button and prompt:
            with st.spinner("Generating image..."):
                # Create mask from rectangle coordinates
                mask = Image.new("L", image.size, 0)
                draw = ImageDraw.Draw(mask)
                draw.rectangle([(x1, y1), (x2, y2)], fill=255)
                
                # Calculate dimensions for generation
                width, height = calculate_optimal_dimensions(image)
                
                # Progress bar
                progress_bar = st.progress(0)
                
                # Generate the image
                try:
                    # Set up progress bar updates
                    progress_text = st.empty()
                    debug_info = st.empty()
                    
                    # Show parameters for debugging
                    debug_info.info(f"Model type: {pipe.__class__.__name__}")
                    
                    # Update progress
                    progress_bar.progress(0.1)
                    progress_text.text("Preparing image and mask...")
                    
                    # Make sure mask is in the right format
                    # Some models require masks where white (255) is the area to inpaint
                    mask_img = mask.convert("L")
                    
                    # Prepare arguments - different models may have different parameter names
                    model_class_name = pipe.__class__.__name__
                    
                    # Common parameters for all models
                    common_params = {
                        "prompt": prompt,
                        "image": image,
                        "mask_image": mask_img,
                        "num_inference_steps": num_inference_steps,
                        "generator": torch.Generator("cpu").manual_seed(seed)
                    }
                    
                    # Add parameters for Flux model
                    common_params["guidance_scale"] = guidance_scale
                    
                    # Try running generation with dimensions
                    try:
                        progress_text.text("Running generation...")
                        progress_bar.progress(0.2)
                        
                        # First try with dimensions
                        common_params["height"] = int(height)
                        common_params["width"] = int(width)
                        result = pipe(**common_params)
                    except Exception as e:
                        debug_info.warning(f"First attempt failed: {str(e)}")
                        progress_text.text("Retrying with adjusted parameters...")
                        
                        # Remove dimensions and try again
                        del common_params["height"]
                        del common_params["width"]
                        result = pipe(**common_params)
                    
                    # Get the result image
                    result_image = result.images[0]
                    
                    # Update final progress
                    progress_bar.progress(1.0)
                    progress_text.text("Complete!")
                    debug_info.empty()  # Clear debug info
                    
                    # Display the result
                    st.image(result_image, caption="Generated Result", use_column_width=True)
                    
                    # Add download button
                    buf = io.BytesIO()
                    result_image.save(buf, format="PNG")
                    st.download_button(
                        label="Download result",
                        data=buf.getvalue(),
                        file_name="inpaint_result.png",
                        mime="image/png",
                    )
                    
                    # Display used seed
                    st.write(f"Seed used: {seed}")
                    
                except Exception as e:
                    st.error(f"An error occurred during generation: {str(e)}")
                    st.error("Try adjusting the parameters or using a different image.")

# If no image is uploaded
if uploaded_file is None:
    with col2:
        st.write("Please upload an image first")