Create app.py

app.py

@@ -0,0 +1,202 @@
+import streamlit as st
+import numpy as np
+import torch
+import random
+from diffusers import FluxFillPipeline
+from PIL import Image
+import io
+
+# 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
+@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):
+    # 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_column_width=True)
+        
+        # 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="white",
+            stroke_width=10,
+            stroke_color="white",
+            background_color="transparent",
+            background_image=image,
+            update_streamlit=True,
+            height=600,
+            drawing_mode="freedraw",
+            key="canvas",
+        )
+        
+        # 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 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")