revert to previous

app.py

@@ -1,4 +1,5 @@
 import spaces
+
 import gradio as gr
 import numpy as np
 import os
@@ -9,20 +10,18 @@ import torch
 from torchvision import transforms
 import zipfile
 
-from diffusers import FluxFillPipeline, AutoencoderKL, DDIMScheduler, DDPMScheduler
-from diffusers.models.attention_processor import AttnProcessor2_0
+from diffusers import FluxFillPipeline, AutoencoderKL
 from PIL import Image
+# from samgeo.text_sam import LangSAM
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 2048
 
-# Initialize the pipeline
-pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16).to("cuda")
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# sam = LangSAM(model_type="sam2-hiera-large").to(device)
 
-# Enable memory efficient attention
-pipe.enable_xformers_memory_efficient_attention()
+pipe = FluxFillPipeline.from_pretrained("black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16).to("cuda")
 
-# Load LoRA models from JSON
 with open("lora_models.json", "r") as f:
     lora_models = json.load(f)
 
@@ -40,7 +39,7 @@ for model_name, model_path in lora_models.items():
 
 lora_models["None"] = None
 
-def calculate_optimal_dimensions(image: Image.Image, scale_factor=1.0):
+def calculate_optimal_dimensions(image: Image.Image):
     # Extract the original dimensions
     original_width, original_height = image.size
 
@@ -60,10 +59,6 @@ def calculate_optimal_dimensions(image: Image.Image, scale_factor=1.0):
         height = FIXED_DIMENSION
         width = round(FIXED_DIMENSION * original_aspect_ratio)
 
-    # Apply scaling factor
-    width = round(width * scale_factor)
-    height = round(height * scale_factor)
-
     # Ensure dimensions are multiples of 8
     width = (width // 8) * 8
     height = (height // 8) * 8
@@ -76,116 +71,21 @@ def calculate_optimal_dimensions(image: Image.Image, scale_factor=1.0):
         height = (width / MIN_ASPECT_RATIO // 8) * 8
 
     # Ensure width and height remain above the minimum dimensions
-    width = max(width, 576)
-    height = max(height, 576)
-
-    # Ensure dimensions don't exceed maximum
-    width = min(width, MAX_IMAGE_SIZE)
-    height = min(height, MAX_IMAGE_SIZE)
+    width = max(width, 576) if width == FIXED_DIMENSION else width
+    height = max(height, 576) if height == FIXED_DIMENSION else height
 
     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, 
-          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
+def infer(edit_images, prompt, lora_model, strength, seed=42, randomize_seed=False, guidance_scale=3.5, num_inference_steps=28, progress=gr.Progress(track_tqdm=True)):
+    # pipe.enable_xformers_memory_efficient_attention()
+    gr.Info("Infering")
+
     if lora_model != "None":
         pipe.load_lora_weights(lora_models[lora_model])
         pipe.enable_lora()
-    else:
-        pipe.disable_lora()
 
-    gr.Info("Processing input images")
+    gr.Info("starting checks")
 
     image = edit_images["background"]
     mask = edit_images["layers"][0]
@@ -194,82 +94,34 @@ def infer(edit_images, prompt, lora_model, strength, seed=42, randomize_seed=Fal
         gr.Info("Please upload an image.")
         return None, None
 
-    # Calculate dimensions with scale factor
-    width, height = calculate_optimal_dimensions(image, scale_factor)
-    
+
+    width, height = calculate_optimal_dimensions(image)
     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}")
-    
-    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
-    
-    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:
-            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)
-    
-    output_image_jpg = result_image.convert("RGB")
+
+    # controlImage = processor(image)
+    gr.Info("generating image")
+    image = pipe(
+        # mask_image_latent=vae.encode(controlImage),
+        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=torch.Generator(device='cuda').manual_seed(seed),
+        # generator=torch.Generator().manual_seed(seed),
+        # lora_scale=0.75 // not supported in this version
+    ).images[0]
+
+    output_image_jpg = image.convert("RGB")
     output_image_jpg.save("output.jpg", "JPEG")
 
     return output_image_jpg, seed
+    # return image, seed
 
 def download_image(image):
     if isinstance(image, np.ndarray):
@@ -277,10 +129,7 @@ def download_image(image):
     image.save("output.png", "PNG")
     return "output.png"
 
-def save_details(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):
+def save_details(result, edit_image, prompt, lora_model, strength, seed, guidance_scale, num_inference_steps):
     image = edit_image["background"]
     mask = edit_image["layers"][0]
 
@@ -298,21 +147,10 @@ def save_details(result, edit_image, prompt, lora_model, strength, seed, guidanc
     details = {
         "prompt": prompt,
         "lora_model": lora_model,
-        "lora_scale": lora_scale,
         "strength": strength,
         "seed": seed,
         "guidance_scale": guidance_scale,
-        "num_inference_steps": num_inference_steps,
-        "scale_factor": scale_factor,
-        "width": result.width,
-        "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
+        "num_inference_steps": num_inference_steps
     }
 
     with open("details.json", "w") as f:
@@ -330,11 +168,16 @@ 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)
+# def generate_mask(image, click_x, click_y):
+#     text_prompt = "face"
+#     mask = sam.predict(image, text_prompt, box_threshold=0.24, text_threshold=0.24)
+#     return mask
 
 examples = [
-    "photography of a young woman, accent lighting, (front view:1.4), detailed skin texture, 8k",
+    "photography of a young woman,  accent lighting,  (front view:1.4),  "
+    # "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",
 ]
 
 css="""
@@ -342,20 +185,12 @@ 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] Advanced Inpainting Tool
-        *Now with layered diffusion for improved quality and control*
+        gr.Markdown(f"""# FLUX.1 [dev]
         """)
         with gr.Row():
             with gr.Column():
@@ -366,13 +201,14 @@ with gr.Blocks(css=css) as demo:
                     image_mode='RGB',
                     layers=False,
                     brush=gr.Brush(colors=["#FFFFFF"]),
+                    # height=600
                 )
                 prompt = gr.Text(
                     label="Prompt",
-                    show_label=True,
+                    show_label=False,
                     max_lines=2,
                     placeholder="Enter your prompt",
-                    container=True,
+                    container=False,
                 )
 
                 lora_model = gr.Dropdown(
@@ -385,148 +221,67 @@ 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(
-                    label="Seed",
-                    minimum=0,
-                    maximum=MAX_SEED,
-                    step=1,
-                    value=0,
-                )
-                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=0,
+            )
+
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
             with gr.Row():
+
                 guidance_scale = gr.Slider(
-                    label="Guidance Scale (Standard Mode)",
+                    label="Guidance Scale",
                     minimum=1,
                     maximum=30,
                     step=0.5,
-                    value=3.5,
+                    value=50,
                 )
 
                 num_inference_steps = gr.Slider(
-                    label="Inference Steps (Standard Mode)",
+                    label="Number of inference steps",
                     minimum=1,
                     maximum=50,
                     step=1,
                     value=28,
                 )
-            
+
             with gr.Row():
-                blur_mask = gr.Slider(
-                    label="Mask Blur Radius",
+
+                strength = gr.Slider(
+                    label="Strength",
                     minimum=0,
-                    maximum=30,
-                    step=1,
-                    value=8,
-                    info="Higher values create smoother transitions at mask boundaries"
-                )
-                
-                detail_level = gr.Slider(
-                    label="Detail Enhancement",
-                    minimum=0.5,
-                    maximum=2.0,
-                    step=0.1,
-                    value=1.2,
-                    info="Controls the sharpness of the final image"
+                    maximum=1,
+                    step=0.01,
+                    value=0.85,
                 )
-            
-            # 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]
-            )
+
+            #     width = gr.Slider(
+            #         label="width",
+            #         minimum=512,
+            #         maximum=3072,
+            #         step=1,
+            #         value=1024,
+            #     )
+
+            #     height = gr.Slider(
+            #         label="height",
+            #         minimum=512,
+            #         maximum=3072,
+            #         step=1,
+            #         value=1024,
+            #     )
 
     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,
-            use_layered_diffusion, blur_mask, detail_level,
-            structure_guidance, texture_guidance, structure_steps, texture_steps
-        ],
-        outputs=[result, seed]
+        fn = infer,
+        inputs = [edit_image, prompt, lora_model, strength, seed, randomize_seed, guidance_scale, num_inference_steps],
+        outputs = [result, seed]
     )
 
     download_button = gr.Button("Download Image as PNG")
@@ -547,25 +302,26 @@ 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,
-                use_layered_diffusion, blur_mask, detail_level,
-                structure_guidance, texture_guidance, structure_steps, texture_steps
-            ],
+            inputs=[result, edit_image, prompt, lora_model, strength, seed, guidance_scale, num_inference_steps],
             outputs=gr.File(label="Download/Save Status")
     )
 
-# Get authentication credentials
+    # edit_image.select(
+    #     fn=generate_mask,
+    #     inputs=[edit_image, gr.Number(), gr.Number()],
+    #     outputs=[edit_image]
+    # )
+
+# demo.launch()
 PASSWORD = os.getenv("GRADIO_PASSWORD")
 USERNAME = os.getenv("GRADIO_USERNAME")
-
-# Create an authentication function
+# Create an authentication object
 def authenticate(username, password):
     if username == USERNAME and password == PASSWORD:
         return True
+
     else:
         return False
-
 # Launch the app with authentication
+
 demo.launch(debug=True, auth=authenticate)