Update app.py

app.py

@@ -1,3 +1,88 @@
+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
+
+# 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"""
@@ -33,4 +118,251 @@ def load_model():
             
             Note: You can find your token at https://huggingface.co/settings/tokens
             """)
-        st.stop()
+        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")