Update app.py

app.py

@@ -1,5 +1,4 @@
 import spaces
-
 import gradio as gr
 import numpy as np
 import os
@@ -12,16 +11,14 @@ import zipfile
 
 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
 
-# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# sam = LangSAM(model_type="sam2-hiera-large").to(device)
-
+# Initialize the pipeline
 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)
 
@@ -39,7 +36,7 @@ for model_name, model_path in lora_models.items():
 
 lora_models["None"] = None
 
-def calculate_optimal_dimensions(image: Image.Image):
+def calculate_optimal_dimensions(image: Image.Image, scale_factor=1.0):
     # Extract the original dimensions
     original_width, original_height = image.size
 
@@ -59,6 +56,10 @@ def calculate_optimal_dimensions(image: Image.Image):
         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
@@ -71,19 +72,27 @@ def calculate_optimal_dimensions(image: Image.Image):
         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
+    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)
 
     return width, height
 
 @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, progress=gr.Progress(track_tqdm=True)):
-    # pipe.enable_xformers_memory_efficient_attention()
+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")
 
+    # Load and enable LoRA if selected
     if lora_model != "None":
         pipe.load_lora_weights(lora_models[lora_model])
         pipe.enable_lora()
+    else:
+        pipe.disable_lora()
 
     gr.Info("starting checks")
 
@@ -94,34 +103,50 @@ def infer(edit_images, prompt, lora_model, strength, seed=42, randomize_seed=Fal
         gr.Info("Please upload an image.")
         return None, None
 
-
-    width, height = calculate_optimal_dimensions(image)
+    # Calculate dimensions with scale factor
+    width, height = calculate_optimal_dimensions(image, scale_factor)
+    
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
-    # 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")
+    # 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,
+        "negative_prompt": negative_prompt,
+        "image": image,
+        "mask_image": 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 in this model version - 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
+
+    output_image_jpg = result_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):
@@ -129,7 +154,8 @@ 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):
+def save_details(result, edit_image, prompt, lora_model, strength, seed, guidance_scale, 
+                num_inference_steps, lora_scale, scale_factor):
     image = edit_image["background"]
     mask = edit_image["layers"][0]
 
@@ -146,11 +172,16 @@ 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,
         "seed": seed,
         "guidance_scale": guidance_scale,
-        "num_inference_steps": num_inference_steps
+        "num_inference_steps": num_inference_steps,
+        "scale_factor": scale_factor,
+        "width": result.width,
+        "height": result.height
     }
 
     with open("details.json", "w") as f:
@@ -168,16 +199,8 @@ def save_details(result, edit_image, prompt, lora_model, strength, seed, guidanc
 def set_image_as_inpaint(image):
     return image
 
-# 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),  "
-    # "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",
+    "photography of a young woman, accent lighting, (front view:1.4)",
 ]
 
 css="""
@@ -190,8 +213,7 @@ css="""
 with gr.Blocks(css=css) as demo:
 
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(f"""# FLUX.1 [dev]
-        """)
+        gr.Markdown(f"""# FLUX.1 [dev] Inpainting Tool""")
         with gr.Row():
             with gr.Column():
                 edit_image = gr.ImageEditor(
@@ -201,16 +223,15 @@ 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=False,
+                    show_label=True,
                     max_lines=2,
                     placeholder="Enter your prompt",
-                    container=False,
+                    container=True,
                 )
-
+                
                 lora_model = gr.Dropdown(
                     label="Select LoRA Model",
                     choices=list(lora_models.keys()),
@@ -222,25 +243,23 @@ with gr.Blocks(css=css) as demo:
             result = gr.Image(label="Result", show_label=False)
 
         with gr.Accordion("Advanced Settings", open=False):
-
-            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():
+                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",
                     minimum=1,
                     maximum=30,
                     step=0.5,
-                    value=50,
+                    value=3.5,
                 )
 
                 num_inference_steps = gr.Slider(
@@ -252,7 +271,6 @@ with gr.Blocks(css=css) as demo:
                 )
 
             with gr.Row():
-
                 strength = gr.Slider(
                     label="Strength",
                     minimum=0,
@@ -260,28 +278,32 @@ with gr.Blocks(css=css) as demo:
                     step=0.01,
                     value=0.85,
                 )
-
-            #     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,
-            #     )
+                
+                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)"
+                )
 
     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],
-        outputs = [result, seed]
+        fn=infer,
+        inputs=[edit_image, prompt, lora_model, strength, seed, randomize_seed, 
+                guidance_scale, num_inference_steps, lora_scale, scale_factor],
+        outputs=[result, seed]
     )
 
     download_button = gr.Button("Download Image as PNG")
@@ -302,26 +324,21 @@ 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],
+            inputs=[result, edit_image, prompt, lora_model, strength, seed, guidance_scale, 
+                   num_inference_steps, lora_scale, scale_factor],
             outputs=gr.File(label="Download/Save Status")
     )
 
-    # edit_image.select(
-    #     fn=generate_mask,
-    #     inputs=[edit_image, gr.Number(), gr.Number()],
-    #     outputs=[edit_image]
-    # )
-
-# demo.launch()
+# Get authentication credentials
 PASSWORD = os.getenv("GRADIO_PASSWORD")
 USERNAME = os.getenv("GRADIO_USERNAME")
-# Create an authentication object
+
+# Create an authentication function
 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)
+# Launch the app with authentication
+demo.launch(debug=True, auth=authenticate)