Update app.py

app.py

@@ -4,51 +4,27 @@ import torch
 import random
 from diffusers import FluxFillPipeline
 from PIL import Image
-import io
+from io import BytesIO
+import base64
+
+# Set page configuration
+st.set_page_config(page_title="FLUX.1 Fill [dev]", layout="wide")
 
 # 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)]
-""")
-
-# Get Hugging Face token
-hf_token = st.text_input("Enter your Hugging Face token (needed to access FLUX.1-Fill-dev)", type="password")
-if not hf_token:
-    st.warning("You need to provide your Hugging Face token to access this model")
-    st.markdown("1. Sign up/login at [Hugging Face](https://huggingface.co/)")
-    st.markdown("2. Generate a token at https://huggingface.co/settings/tokens")
-    st.markdown("3. Request access to [FLUX.1-Fill-dev](https://huggingface.co/black-forest-labs/FLUX.1-Fill-dev)")
-    st.stop()
-
-# Load the model
+# Initialize the model
 @st.cache_resource
-def load_model(token):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    try:
-        return FluxFillPipeline.from_pretrained(
-            "black-forest-labs/FLUX.1-Fill-dev", 
-            torch_dtype=torch.bfloat16,
-            use_auth_token=token
-        ).to(device)
-    except Exception as e:
-        st.error(f"Error loading model: {str(e)}")
-        st.stop()
-
-pipe = load_model(hf_token)
-
-def calculate_optimal_dimensions(image: Image.Image):
+def load_model():
+    pipe = FluxFillPipeline.from_pretrained(
+        "black-forest-labs/FLUX.1-Fill-dev", 
+        torch_dtype=torch.bfloat16
+    ).to("cuda")
+    return pipe
+
+# Function to calculate optimal dimensions
+def calculate_optimal_dimensions(image):
     # Extract the original dimensions
     original_width, original_height = image.size
     
@@ -83,120 +59,154 @@ def calculate_optimal_dimensions(image: Image.Image):
     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])
+    return int(width), int(height)
 
-with col1:
-    # Upload image
-    uploaded_file = st.file_uploader("Upload an image for inpainting", type=["jpg", "jpeg", "png"])
+# Custom component for image editing and masking
+def image_editor():
+    uploaded_file = st.file_uploader("Upload an image for inpainting", type=["png", "jpg", "jpeg"])
     
     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)
+        width, height = image.size
+        
+        # Create a placeholder for the canvas
+        canvas_placeholder = st.empty()
+        
+        # Display instructions
+        st.markdown("### Draw a mask on the image")
+        st.markdown("Draw on the areas you want to edit (white areas will be inpainted)")
         
-        # Canvas for creating mask
-        st.write("Draw on the image to create a mask for inpainting")
+        # Create canvas for mask drawing
+        # Note: We need to import the components within function to avoid import errors
+        import streamlit_drawable_canvas
         from streamlit_drawable_canvas import st_canvas
+        
         canvas_result = st_canvas(
-            fill_color="white",
+            fill_color="rgba(255, 255, 255, 0.3)",
             stroke_width=10,
-            stroke_color="white",
-            background_color="transparent",
+            stroke_color="#FFFFFF",
             background_image=image,
-            update_streamlit=True,
-            height=600,
+            height=min(600, height),
+            width=min(1000, width),
             drawing_mode="freedraw",
             key="canvas",
         )
         
+        if canvas_result.image_data is not None:
+            # Extract mask from canvas
+            mask_data = canvas_result.image_data
+            mask = Image.fromarray((mask_data[:, :, 3] > 0).astype(np.uint8) * 255)
+            
+            return {
+                "background": image,
+                "mask": mask,
+                "has_mask": (mask.getextrema()[1] > 0)
+            }
+    
+    return None
+
+# Function to run inference
+def run_inference(image_data, prompt, seed, randomize_seed, guidance_scale, num_inference_steps):
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+        
+    pipe = load_model()
+    width, height = calculate_optimal_dimensions(image_data["background"])
+    
+    with st.spinner("Generating image..."):
+        generated_image = pipe(
+            prompt=prompt,
+            image=image_data["background"],
+            mask_image=image_data["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 generated_image, seed
+
+# 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",
+]
+
+# Main UI
+def main():
+    # Header
+    st.title("FLUX.1 Fill [dev]")
+    st.markdown("""
+    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)]
+    """)
+    
+    # Create columns for layout
+    col1, col2 = st.columns([1, 1])
+    
+    with col1:
+        # Image editor
+        image_data = image_editor()
+        
         # 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
+        st.subheader("Example Prompts")
+        for i, example in enumerate(examples):
+            if st.button(f"Use Example: {example}", key=f"example_{i}"):
+                st.session_state.prompt = example
+                st.experimental_rerun()
+    
+        # Advanced settings
         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)
-        
+            num_inference_steps = st.slider("Number of inference steps", 1, 50, 28, 1)
+    
         # 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 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)
+        run_button = st.button("Run")
+    
+    # Result display in second column
+    with col2:
+        # Check if we need to run inference
+        if image_data is not None and prompt and run_button:
+            if not image_data.get("has_mask", False):
+                st.error("Please draw a mask on the image")
+            else:
+                result_image, result_seed = run_inference(
+                    image_data, 
+                    prompt, 
+                    seed, 
+                    randomize_seed, 
+                    guidance_scale, 
+                    num_inference_steps
+                )
                 
-                # Progress bar
-                progress_bar = st.progress(0)
+                # Update seed if it was randomized
+                if randomize_seed:
+                    st.session_state.seed = result_seed
                 
-                # Generate the image
-                def update_progress(step, total_steps):
-                    progress_bar.progress(step / total_steps)
+                # Display result
+                st.subheader("Generated Image")
+                st.image(result_image, use_column_width=True)
                 
-                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")
+                # Download button
+                buf = BytesIO()
+                result_image.save(buf, format="PNG")
+                byte_im = buf.getvalue()
+                download_button_str = f"""
+                <a href="data:image/png;base64,{base64.b64encode(byte_im).decode()}" download="generated_image.png">
+                    <div style="display: inline-flex; align-items: center; background-color: #4CAF50; color: white; padding: 0.5em 1em; border-radius: 4px; text-decoration: none; font-weight: bold; margin-top: 10px;">
+                        ⬇️ Download Image
+                    </div>
+                </a>
+                """
+                st.markdown(download_button_str, unsafe_allow_html=True)
+
+if __name__ == "__main__":
+    main()