Update app.py

app.py

@@ -9,7 +9,8 @@ import torch
 from torchvision import transforms
 import zipfile
 
-from diffusers import FluxFillPipeline, AutoencoderKL
+from diffusers import FluxFillPipeline, AutoencoderKL, DDIMScheduler, DDPMScheduler
+from diffusers.models.attention_processor import AttnProcessor2_0
 from PIL import Image
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -18,6 +19,9 @@ MAX_IMAGE_SIZE = 2048
 # Initialize the pipeline
 pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16).to("cuda")
 
+# Enable memory efficient attention
+pipe.enable_xformers_memory_efficient_attention()
+
 # Load LoRA models from JSON
 with open("lora_models.json", "r") as f:
     lora_models = json.load(f)
@@ -81,11 +85,98 @@ def calculate_optimal_dimensions(image: Image.Image, scale_factor=1.0):
 
     return width, height
 
+def preprocess_mask(mask, blur_radius=10):
+    """Apply blurring to create a soft mask for smoother transitions"""
+    from PIL import ImageFilter
+    if mask:
+        # Apply Gaussian blur to soften the mask edges
+        blurred_mask = mask.filter(ImageFilter.GaussianBlur(radius=blur_radius))
+        return blurred_mask
+    return mask
+
+def enhance_detail(img, sharpness_factor=1.5):
+    """Enhance the details in the generated image"""
+    from PIL import ImageEnhance
+    if img:
+        enhancer = ImageEnhance.Sharpness(img)
+        enhanced_img = enhancer.enhance(sharpness_factor)
+        return enhanced_img
+    return img
+
+def layer_based_inference(pipe, image, mask, prompt, 
+                          structure_guidance_scale, texture_guidance_scale, 
+                          structure_steps, texture_steps, strength, seed,
+                          lora_scale, width, height):
+    """Perform a two-stage layer-based diffusion process for better quality"""
+    # Create generators with the same seed for reproducibility
+    structure_generator = torch.Generator(device='cuda').manual_seed(seed)
+    texture_generator = torch.Generator(device='cuda').manual_seed(seed+1)
+    
+    # Configure structure pipeline (focus on shapes and composition)
+    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
+    
+    # Stage 1: Generate the overall structure with more steps but lower guidance
+    structure_kwargs = {
+        "prompt": f"structure of {prompt}",
+        "image": image,
+        "mask_image": mask,
+        "height": height,
+        "width": width,
+        "guidance_scale": structure_guidance_scale,
+        "strength": strength * 0.9,  # Less strength to preserve original structure
+        "num_inference_steps": structure_steps,
+        "generator": structure_generator,
+    }
+    
+    # Add LoRA scale if supported
+    if pipe.is_lora_enabled():
+        structure_kwargs["cross_attention_kwargs"] = {"scale": lora_scale * 0.8}  # Reduce LoRA impact for structure
+    
+    # Generate the structural base
+    try:
+        intermediate_result = pipe(**structure_kwargs).images[0]
+    except Exception as e:
+        print(f"Structure generation error: {str(e)}")
+        return None
+    
+    # Configure texture pipeline (focus on details and textures)
+    pipe.scheduler = DDPMScheduler.from_config(pipe.scheduler.config)
+    
+    # Stage 2: Refine with texture details using the intermediate result
+    texture_kwargs = {
+        "prompt": f"detailed texture of {prompt}",
+        "image": intermediate_result,
+        "mask_image": mask,
+        "height": height,
+        "width": width,
+        "guidance_scale": texture_guidance_scale,
+        "strength": strength * 0.6,  # Lower strength to maintain structure
+        "num_inference_steps": texture_steps,
+        "generator": texture_generator,
+    }
+    
+    # Add LoRA scale if supported
+    if pipe.is_lora_enabled():
+        texture_kwargs["cross_attention_kwargs"] = {"scale": lora_scale * 1.2}  # Increase LoRA impact for texture
+    
+    # Generate the final result with detailed textures
+    try:
+        final_result = pipe(**texture_kwargs).images[0]
+        return final_result
+    except Exception as e:
+        print(f"Texture generation error: {str(e)}")
+        return intermediate_result  # Return intermediate result if texture stage fails
+
 @spaces.GPU(durations=300)
 def infer(edit_images, prompt, lora_model, strength, seed=42, randomize_seed=False, 
           guidance_scale=3.5, num_inference_steps=28, lora_scale=0.75, 
-          scale_factor=1.0, progress=gr.Progress(track_tqdm=True)):
-    gr.Info("Infering")
+          scale_factor=1.0, 
+          use_layered_diffusion=True, blur_mask=8, detail_level=1.2,
+          structure_guidance=2.5, texture_guidance=5.0,
+          structure_steps=20, texture_steps=15,
+          progress=gr.Progress(track_tqdm=True)):
+    
+    gr.Info("Starting inference process")
 
     # Load and enable LoRA if selected
     if lora_model != "None":
@@ -94,7 +185,7 @@ def infer(edit_images, prompt, lora_model, strength, seed=42, randomize_seed=Fal
     else:
         pipe.disable_lora()
 
-    gr.Info("starting checks")
+    gr.Info("Processing input images")
 
     image = edit_images["background"]
     mask = edit_images["layers"][0]
@@ -108,40 +199,73 @@ def infer(edit_images, prompt, lora_model, strength, seed=42, randomize_seed=Fal
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
-
+    
+    # Process the mask for smoother transitions
+    processed_mask = preprocess_mask(mask, blur_radius=blur_mask)
+    
     # Generate image
     gr.Info(f"Generating image at {width}x{height}")
-    generator = torch.Generator(device='cuda').manual_seed(seed)
     
-    # Configure pipeline parameters
-    pipeline_kwargs = {
-        "prompt": prompt,
-        "prompt_2": prompt,
-        "image": image,
-        "mask_image": mask,
-        "height": height,
-        "width": width,
-        "guidance_scale": guidance_scale,
-        "strength": strength,
-        "num_inference_steps": num_inference_steps,
-        "generator": generator,
-    }
+    if use_layered_diffusion:
+        gr.Info("Using layered diffusion for higher quality output")
+        result_image = layer_based_inference(
+            pipe=pipe,
+            image=image,
+            mask=processed_mask,
+            prompt=prompt,
+            structure_guidance_scale=structure_guidance,
+            texture_guidance_scale=texture_guidance,
+            structure_steps=structure_steps,
+            texture_steps=texture_steps,
+            strength=strength,
+            seed=seed,
+            lora_scale=lora_scale,
+            width=width,
+            height=height
+        )
+        
+        if result_image is None:
+            gr.Error("Layered diffusion failed. Falling back to standard diffusion.")
+            use_layered_diffusion = False
     
-    # Add LoRA scale if model supports it
-    if lora_model != "None":
+    if not use_layered_diffusion:
+        # Standard diffusion as fallback
+        generator = torch.Generator(device='cuda').manual_seed(seed)
+        
+        # Configure pipeline parameters
+        pipeline_kwargs = {
+            "prompt": prompt,
+            "prompt_2": prompt,
+            "image": image,
+            "mask_image": processed_mask,
+            "height": height,
+            "width": width,
+            "guidance_scale": guidance_scale,
+            "strength": strength,
+            "num_inference_steps": num_inference_steps,
+            "generator": generator,
+        }
+        
+        # Add LoRA scale if model supports it
+        if lora_model != "None":
+            try:
+                pipeline_kwargs["cross_attention_kwargs"] = {"scale": lora_scale}
+            except:
+                gr.Info("LoRA scale not supported - using default scaling")
+        
+        # Run the pipeline
         try:
-            pipeline_kwargs["cross_attention_kwargs"] = {"scale": lora_scale}
-        except:
-            gr.Info("LoRA scale not supported in this model version - using default scaling")
+            output = pipe(**pipeline_kwargs)
+            result_image = output.images[0]
+        except Exception as e:
+            gr.Error(f"Error during generation: {str(e)}")
+            return None, seed
+    
+    # Enhance details based on user preference
+    if detail_level > 1.0:
+        gr.Info("Enhancing image details")
+        result_image = enhance_detail(result_image, sharpness_factor=detail_level)
     
-    # Run the pipeline
-    try:
-        output = pipe(**pipeline_kwargs)
-        result_image = output.images[0]
-    except Exception as e:
-        gr.Error(f"Error during generation: {str(e)}")
-        return None, seed
-
     output_image_jpg = result_image.convert("RGB")
     output_image_jpg.save("output.jpg", "JPEG")
 
@@ -154,7 +278,9 @@ def download_image(image):
     return "output.png"
 
 def save_details(result, edit_image, prompt, lora_model, strength, seed, guidance_scale, 
-                num_inference_steps, lora_scale, scale_factor):
+                num_inference_steps, lora_scale, scale_factor,
+                use_layered_diffusion, blur_mask, detail_level,
+                structure_guidance, texture_guidance, structure_steps, texture_steps):
     image = edit_image["background"]
     mask = edit_image["layers"][0]
 
@@ -171,7 +297,6 @@ def save_details(result, edit_image, prompt, lora_model, strength, seed, guidanc
 
     details = {
         "prompt": prompt,
-        "negative_prompt": negative_prompt,
         "lora_model": lora_model,
         "lora_scale": lora_scale,
         "strength": strength,
@@ -180,7 +305,14 @@ def save_details(result, edit_image, prompt, lora_model, strength, seed, guidanc
         "num_inference_steps": num_inference_steps,
         "scale_factor": scale_factor,
         "width": result.width,
-        "height": result.height
+        "height": result.height,
+        "use_layered_diffusion": use_layered_diffusion,
+        "blur_mask": blur_mask,
+        "detail_level": detail_level,
+        "structure_guidance": structure_guidance,
+        "texture_guidance": texture_guidance,
+        "structure_steps": structure_steps,
+        "texture_steps": texture_steps
     }
 
     with open("details.json", "w") as f:
@@ -198,8 +330,11 @@ def save_details(result, edit_image, prompt, lora_model, strength, seed, guidanc
 def set_image_as_inpaint(image):
     return image
 
+def toggle_layered_diffusion(use_layers):
+    return gr.update(visible=use_layers)
+
 examples = [
-    "photography of a young woman, accent lighting, (front view:1.4)",
+    "photography of a young woman, accent lighting, (front view:1.4), detailed skin texture, 8k",
 ]
 
 css="""
@@ -207,12 +342,21 @@ css="""
     margin: 0 auto;
     max-width: 1000px;
 }
+.layer-settings {
+    border: 1px solid #ccc;
+    padding: 10px;
+    border-radius: 8px;
+    background-color: #f9f9f9;
+    margin-top: 10px;
+}
 """
 
 with gr.Blocks(css=css) as demo:
 
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""# FLUX.1 [dev] Inpainting Tool""")
+        gr.Markdown(f"""# FLUX.1 [dev] Advanced Inpainting Tool
+        *Now with layered diffusion for improved quality and control*
+        """)
         with gr.Row():
             with gr.Column():
                 edit_image = gr.ImageEditor(
@@ -230,7 +374,7 @@ with gr.Blocks(css=css) as demo:
                     placeholder="Enter your prompt",
                     container=True,
                 )
-                
+
                 lora_model = gr.Dropdown(
                     label="Select LoRA Model",
                     choices=list(lora_models.keys()),
@@ -241,6 +385,42 @@ with gr.Blocks(css=css) as demo:
 
             result = gr.Image(label="Result", show_label=False)
 
+        with gr.Accordion("Basic Settings", open=True):
+            with gr.Row():
+                strength = gr.Slider(
+                    label="Strength",
+                    minimum=0,
+                    maximum=1,
+                    step=0.01,
+                    value=0.85,
+                    info="Controls how much to modify the original image"
+                )
+                
+                lora_scale = gr.Slider(
+                    label="LoRA Scale",
+                    minimum=0,
+                    maximum=2,
+                    step=0.05,
+                    value=0.75,
+                    info="Controls the influence of the LoRA model"
+                )
+            
+            with gr.Row():
+                scale_factor = gr.Slider(
+                    label="Image Scale Factor",
+                    minimum=0.5,
+                    maximum=2.0,
+                    step=0.1,
+                    value=1.0,
+                    info="Scale factor for image dimensions (1.0 = original, 2.0 = double size)"
+                )
+                
+                use_layered_diffusion = gr.Checkbox(
+                    label="Use Layered Diffusion", 
+                    value=True,
+                    info="Process in structure and texture layers for better quality"
+                )
+
         with gr.Accordion("Advanced Settings", open=False):
             with gr.Row():
                 seed = gr.Slider(
@@ -254,7 +434,7 @@ with gr.Blocks(css=css) as demo:
 
             with gr.Row():
                 guidance_scale = gr.Slider(
-                    label="Guidance Scale",
+                    label="Guidance Scale (Standard Mode)",
                     minimum=1,
                     maximum=30,
                     step=0.5,
@@ -262,46 +442,90 @@ with gr.Blocks(css=css) as demo:
                 )
 
                 num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
+                    label="Inference Steps (Standard Mode)",
                     minimum=1,
                     maximum=50,
                     step=1,
                     value=28,
                 )
-
+            
             with gr.Row():
-                strength = gr.Slider(
-                    label="Strength",
+                blur_mask = gr.Slider(
+                    label="Mask Blur Radius",
                     minimum=0,
-                    maximum=1,
-                    step=0.01,
-                    value=0.85,
+                    maximum=30,
+                    step=1,
+                    value=8,
+                    info="Higher values create smoother transitions at mask boundaries"
                 )
                 
-                lora_scale = gr.Slider(
-                    label="LoRA Scale",
-                    minimum=0,
-                    maximum=2,
-                    step=0.05,
-                    value=0.75,
-                    info="Controls the influence of the LoRA model"
-                )
-            
-            with gr.Row():
-                scale_factor = gr.Slider(
-                    label="Image Scale Factor",
+                detail_level = gr.Slider(
+                    label="Detail Enhancement",
                     minimum=0.5,
                     maximum=2.0,
                     step=0.1,
-                    value=1.0,
-                    info="Scale factor for image dimensions (1.0 = original, 2.0 = double size)"
+                    value=1.2,
+                    info="Controls the sharpness of the final image"
                 )
+            
+            # Layered diffusion settings
+            with gr.Column(visible=True, elem_id="layer-settings", elem_class="layer-settings") as layer_settings:
+                gr.Markdown("### Layered Diffusion Settings")
+                
+                with gr.Row():
+                    structure_guidance = gr.Slider(
+                        label="Structure Guidance Scale",
+                        minimum=1.0,
+                        maximum=10.0,
+                        step=0.1,
+                        value=2.5,
+                        info="Controls adherence to prompt during structure generation (lower = more creative)"
+                    )
+                    
+                    texture_guidance = gr.Slider(
+                        label="Texture Guidance Scale",
+                        minimum=1.0,
+                        maximum=10.0,
+                        step=0.1,
+                        value=5.0,
+                        info="Controls adherence to prompt during texture refinement (higher = more detailed)"
+                    )
+                
+                with gr.Row():
+                    structure_steps = gr.Slider(
+                        label="Structure Steps",
+                        minimum=10,
+                        maximum=40,
+                        step=1,
+                        value=20,
+                        info="Inference steps for structure generation"
+                    )
+                    
+                    texture_steps = gr.Slider(
+                        label="Texture Steps",
+                        minimum=10,
+                        maximum=40,
+                        step=1,
+                        value=15,
+                        info="Inference steps for texture refinement"
+                    )
+            
+            # Toggle visibility of layered settings based on checkbox
+            use_layered_diffusion.change(
+                fn=toggle_layered_diffusion,
+                inputs=[use_layered_diffusion],
+                outputs=[layer_settings]
+            )
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
-        inputs=[edit_image, prompt, lora_model, strength, seed, randomize_seed, 
-                guidance_scale, num_inference_steps, lora_scale, scale_factor],
+        inputs=[
+            edit_image, prompt, lora_model, strength, seed, randomize_seed, 
+            guidance_scale, num_inference_steps, lora_scale, scale_factor,
+            use_layered_diffusion, blur_mask, detail_level,
+            structure_guidance, texture_guidance, structure_steps, texture_steps
+        ],
         outputs=[result, seed]
     )
 
@@ -323,8 +547,12 @@ with gr.Blocks(css=css) as demo:
 
     save_button.click(
             fn=save_details,
-            inputs=[result, edit_image, prompt, lora_model, strength, seed, guidance_scale, 
-                   num_inference_steps, lora_scale, scale_factor],
+            inputs=[
+                result, edit_image, prompt, lora_model, strength, seed, guidance_scale, 
+                num_inference_steps, lora_scale, scale_factor,
+                use_layered_diffusion, blur_mask, detail_level,
+                structure_guidance, texture_guidance, structure_steps, texture_steps
+            ],
             outputs=gr.File(label="Download/Save Status")
     )