Update app.py

app.py

@@ -2,46 +2,66 @@ import streamlit as st
 import numpy as np
 import torch
 import random
-import base64
-from io import BytesIO
 from diffusers import FluxFillPipeline
 from PIL import Image
-from streamlit_drawable_canvas import st_canvas
+import io
 
 # Constants
-MAX_SEED = np.int32().max
+MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-# Function to convert PIL Image to data URL (to replace image_to_url)
-def image_to_base64_url(img):
-    buffered = BytesIO()
-    img.save(buffered, format="PNG")
-    img_str = base64.b64encode(buffered.getvalue()).decode()
-    return f"data:image/png;base64,{img_str}"
-
-
-# Authenticate with Hugging Face
-def authenticate():
-    token = st.text_input("Enter your Hugging Face token:", type="password", 
-                          help="Find your token at https://huggingface.co/settings/tokens")
-    if token:
-        try:
-            login(token)
-            st.success("Successfully authenticated with Hugging Face!")
-            return True
-        except Exception as e:
-            st.error(f"Authentication failed: {e}")
-            return False
-    return False
-    
-# Load model (wrapped in a function to load only when needed)
+# 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. Run `huggingface-cli login` in your terminal (or add your token to Hugging Face Spaces secrets)
+3. Request access to [FLUX.1-Fill-dev](https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev) by clicking 'Access repository'
+4. Wait for approval from model owners
+""")
+
 @st.cache_resource
 def load_model():
-    pipe = FluxFillPipeline.from_pretrained(
-        "black-forest-labs/FLUX.1-Fill-dev", 
-        torch_dtype=torch.bfloat16
-    ).to("cuda")
-    return pipe
+    """Load the model using the Hugging Face CLI login approach"""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    try:
+        # This should work if the user has done huggingface-cli login
+        # or if token is in HF_HOME/.huggingface/token
+        return FluxFillPipeline.from_pretrained(
+            "black-forest-labs/FLUX.1-Fill-dev", 
+            torch_dtype=torch.bfloat16
+        ).to(device)
+    except Exception as e:
+        st.error(f"Error loading model: {str(e)}")
+        if "401 Client Error" in str(e):
+            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. Make sure you've run 'huggingface-cli login' or set up the token in Spaces secrets
+            3. Wait for approval from model owners
+            """)
+        st.stop()
+
+try:
+    pipe = load_model()
+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
@@ -80,119 +100,118 @@ def calculate_optimal_dimensions(image: Image.Image):
 
     return width, height
 
-def infer(pipe, image, mask, prompt, seed, randomize_seed, guidance_scale, num_inference_steps):
-    width, height = calculate_optimal_dimensions(image)
-    
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-    
-    with st.status("Generating image..."):
-        result_image = pipe(
-            prompt=prompt,
-            image=image,
-            mask_image=mask,
-            height=height,
-            width=width,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            generator=torch.Generator("cpu").manual_seed(seed)
-        ).images[0]
-    
-    return result_image, seed
+# Create two columns for layout
+col1, col2 = st.columns([1, 1])
 
-def image_editor():
-    uploaded_image = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])
+with col1:
+    # Upload image
+    uploaded_file = st.file_uploader("Upload an image for inpainting", type=["jpg", "jpeg", "png"])
     
-    if uploaded_image is not None:
-        image = Image.open(uploaded_image).convert("RGB")
-        
-        # Display image for reference
+    if uploaded_file is not None:
+        # Display the uploaded image
+        image = Image.open(uploaded_file).convert("RGB")
         st.image(image, caption="Uploaded Image", use_column_width=True)
         
-        # Create a base64 URL for the background image
-        background_image_url = image_to_base64_url(image)
-        
-        # Create the canvas for drawing masks
+        # Canvas for creating mask
+        st.write("Draw on the image to create a mask for inpainting")
+        from streamlit_drawable_canvas import st_canvas
         canvas_result = st_canvas(
-            fill_color="rgba(255, 255, 255, 0.3)",
+            fill_color="white",
             stroke_width=10,
-            stroke_color="#FFFFFF",
-            background_image=background_image_url,
+            stroke_color="white",
+            background_color="transparent",
+            background_image=image,
+            update_streamlit=True,
             height=600,
             drawing_mode="freedraw",
             key="canvas",
         )
         
-        return {
-            "background": image,
-            "mask": Image.fromarray(canvas_result.image_data[:, :, -1], mode="L") if canvas_result.image_data is not None else None
-        }
-    
-    return None
-
-def main():
-    st.set_page_config(page_title="FLUX.1 Fill", layout="wide")
-    
-    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)]
-    """)
-    
-    # Initialize model
-    pipe = load_model()
-    
-    col1, col2 = st.columns([6, 6])
-    
-    with col1:
-        st.subheader("Input")
-        image_data = image_editor()
-        
-        prompt = st.text_input("Prompt", placeholder="Enter your prompt")
+        # 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",
         ]
         
-        selected_example = st.selectbox("Or try an example prompt:", [""] + examples)
-        if selected_example:
-            prompt = selected_example
-            
+        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"):
-            seed = st.slider("Seed", 0, MAX_SEED, 0)
             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 = st.button("Run")
-    
-    with col2:
-        st.subheader("Result")
-        result_container = st.empty()
-        seed_text = st.empty()
-    
-    if run_button and image_data and prompt:
-        if image_data["mask"] is not None:
-            result_image, used_seed = infer(
-                pipe, 
-                image_data["background"], 
-                image_data["mask"], 
-                prompt, 
-                seed, 
-                randomize_seed, 
-                guidance_scale, 
-                num_inference_steps
-            )
-            
-            result_container.image(result_image, caption="Generated Image", use_column_width=True)
-            seed_text.text(f"Seed used: {used_seed}")
-        else:
-            st.error("Please draw a mask on the image")
+        # Run button
+        run_button = st.button("Generate")
 
-if __name__ == "__main__":
-    main()
+with col2:
+    if uploaded_file is not None:
+        st.write("Result will appear here")
+        
+        if run_button and prompt and canvas_result.image_data is not None:
+            with st.spinner("Generating image..."):
+                # Create mask from canvas
+                mask_data = canvas_result.image_data
+                mask = Image.fromarray(mask_data.astype(np.uint8)).convert("L")
+                
+                # Calculate dimensions
+                width, height = calculate_optimal_dimensions(image)
+                
+                # Progress bar
+                progress_bar = st.progress(0)
+                
+                # Generate the image
+                def update_progress(step, total_steps):
+                    progress_bar.progress(step / total_steps)
+                
+                try:
+                    result_image = pipe(
+                        prompt=prompt,
+                        image=image,
+                        mask_image=mask,
+                        height=int(height),
+                        width=int(width),
+                        guidance_scale=guidance_scale,
+                        num_inference_steps=num_inference_steps,
+                        generator=torch.Generator("cpu").manual_seed(seed),
+                        callback=update_progress
+                    ).images[0]
+                    
+                    # Update final progress
+                    progress_bar.progress(1.0)
+                    
+                    # 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="flux_fill_result.png",
+                        mime="image/png",
+                    )
+                    
+                    # Display used seed
+                    st.write(f"Seed used: {seed}")
+                    
+                except Exception as e:
+                    st.error(f"An error occurred: {str(e)}")
+
+# If no image is uploaded
+if uploaded_file is None:
+    with col2:
+        st.write("Please upload an image first")