ui revamp

app.py

@@ -139,14 +139,43 @@ def upsample_prompt_logic(prompt, image_list):
         print(f"Upsampling failed: {e}")
         return prompt
 
+def update_dimensions_from_image(image_list):
+    """Update width/height sliders based on uploaded image aspect ratio.
+    Keeps one side at 1024 and scales the other proportionally, with both sides as multiples of 8."""
+    if image_list is None or len(image_list) == 0:
+        return 1024, 1024  # Default dimensions
+    
+    # Get the first image to determine dimensions
+    img = image_list[0][0]  # Gallery returns list of tuples (image, caption)
+    img_width, img_height = img.size
+    
+    aspect_ratio = img_width / img_height
+    
+    if aspect_ratio >= 1:  # Landscape or square
+        new_width = 1024
+        new_height = int(1024 / aspect_ratio)
+    else:  # Portrait
+        new_height = 1024
+        new_width = int(1024 * aspect_ratio)
+    
+    # Round to nearest multiple of 8
+    new_width = round(new_width / 8) * 8
+    new_height = round(new_height / 8) * 8
+    
+    # Ensure within valid range (minimum 256, maximum 1024)
+    new_width = max(256, min(1024, new_width))
+    new_height = max(256, min(1024, new_height))
+    
+    return new_width, new_height
+
 # Updated duration function to match generate_image arguments (including progress)
-def get_duration(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, force_dimensions, progress=gr.Progress(track_tqdm=True)):
+def get_duration(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
     num_images = 0 if image_list is None else len(image_list)
     step_duration = 1 + 0.8 * num_images
     return max(65, num_inference_steps * step_duration + 10)
 
 @spaces.GPU(duration=get_duration)
-def generate_image(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, force_dimensions, progress=gr.Progress(track_tqdm=True)):
+def generate_image(prompt_embeds, image_list, width, height, num_inference_steps, guidance_scale, seed, progress=gr.Progress(track_tqdm=True)):
     # Move embeddings to GPU only when inside the GPU decorated function
     prompt_embeds = prompt_embeds.to(device)
     
@@ -158,12 +187,10 @@ def generate_image(prompt_embeds, image_list, width, height, num_inference_steps
         "num_inference_steps": num_inference_steps,
         "guidance_scale": guidance_scale,
         "generator": generator,
+        "width": width,
+        "height": height,
     }
     
-    if image_list is None or force_dimensions:
-        pipe_kwargs["width"] = width
-        pipe_kwargs["height"] = height
-    
     # Progress bar for the actual generation steps
     if progress:
         progress(0, desc="Starting generation...")
@@ -171,7 +198,7 @@ def generate_image(prompt_embeds, image_list, width, height, num_inference_steps
     image = pipe(**pipe_kwargs).images[0]
     return image
 
-def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=50, guidance_scale=2.5, force_dimensions=False, prompt_upsampling=False, progress=gr.Progress(track_tqdm=True)):
+def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024, height=1024, num_inference_steps=50, guidance_scale=2.5, prompt_upsampling=False, progress=gr.Progress(track_tqdm=True)):
     
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -206,7 +233,6 @@ def infer(prompt, input_images=None, seed=42, randomize_seed=False, width=1024,
         num_inference_steps, 
         guidance_scale, 
         seed, 
-        force_dimensions,
         progress
     )
     
@@ -227,7 +253,7 @@ examples_images = [
 css="""
 #col-container {
     margin: 0 auto;
-    max-width: 620px;
+    max-width: 1200px;
 }
 .gallery-container img{
     object-fit: contain;
@@ -240,89 +266,85 @@ with gr.Blocks() as demo:
         gr.Markdown(f"""# FLUX.2 [dev]
 FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and combining images based on text instructions model [[model](https://huggingface.co/black-forest-labs/FLUX.2-dev)], [[blog](https://bfl.ai/blog/flux-2)]
         """)
-
-        with gr.Accordion("Input image(s) (optional)", open=True):
-            input_images = gr.Gallery(
-                label="Input Image(s)",
-                type="pil",
-                columns=3,
-                rows=1,
-            )
-        
         with gr.Row():
-            
-            prompt = gr.Text(
-                label="Prompt",
-                show_label=False,
-                max_lines=2,
-                placeholder="Enter your prompt",
-                container=False,
-                scale=3
-            )
-            
-            run_button = gr.Button("Run", scale=1)
-        
-        result = gr.Image(label="Result", show_label=False)
+            with gr.Column():
+                with gr.Row():
+                    prompt = gr.Text(
+                        label="Prompt",
+                        show_label=False,
+                        max_lines=2,
+                        placeholder="Enter your prompt",
+                        container=False,
+                        scale=3
+                    )
+                    
+                    run_button = gr.Button("Run", scale=1)
+                    
+                with gr.Accordion("Input image(s) (optional)", open=True):
+                    input_images = gr.Gallery(
+                        label="Input Image(s)",
+                        type="pil",
+                        columns=3,
+                        rows=1,
+                    )
+                
+                with gr.Accordion("Advanced Settings", open=False):
+                    prompt_upsampling = gr.Checkbox(
+                        label="Prompt Upsampling",
+                        value=True,
+                        info="Automatically enhance the prompt using a VLM"
+                    )
         
-        with gr.Accordion("Advanced Settings", open=False):
-            
-            prompt_upsampling = gr.Checkbox(
-                label="Prompt Upsampling",
-                value=True,
-                info="Automatically enhance the prompt using a VLM"
-            )
-
-            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():
+                        
+                        width = gr.Slider(
+                            label="Width",
+                            minimum=256,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=8,
+                            value=1024,
+                        )
+                        
+                        height = gr.Slider(
+                            label="Height",
+                            minimum=256,
+                            maximum=MAX_IMAGE_SIZE,
+                            step=8,
+                            value=1024,
+                        )
+                    
+                    with gr.Row():
+                        
+                        num_inference_steps = gr.Slider(
+                            label="Number of inference steps",
+                            minimum=1,
+                            maximum=100,
+                            step=1,
+                            value=30,
+                        )
+                        
+                        guidance_scale = gr.Slider(
+                            label="Guidance scale",
+                            minimum=0.0,
+                            maximum=10.0,
+                            step=0.1,
+                            value=4,
+                        )
                 
-                width = gr.Slider(
-                    label="Width",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
-                )
                 
-                height = gr.Slider(
-                    label="Height",
-                    minimum=256,
-                    maximum=MAX_IMAGE_SIZE,
-                    step=32,
-                    value=1024,
-                )
+            with gr.Column():
+                result = gr.Image(label="Result", show_label=False)
             
-            force_dimensions = gr.Checkbox(
-                label="Force width/height when image input",
-                value=False,
-                info="When unchecked, width/height settings are ignored if input images are provided"
-            )
-            
-            with gr.Row():
-                
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=100,
-                    step=1,
-                    value=30,
-                )
-                
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=4,
-                )
         
         gr.Examples(
             examples=examples,
@@ -342,10 +364,17 @@ FLUX.2 [dev] is a 32B model rectified flow capable of generating, editing and co
             cache_mode="lazy"
         )
 
+    # Auto-update dimensions when images are uploaded
+    input_images.upload(
+        fn=update_dimensions_from_image,
+        inputs=[input_images],
+        outputs=[width, height]
+    )
+
     gr.on(
         triggers=[run_button.click, prompt.submit],
         fn=infer,
-        inputs=[prompt, input_images, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, force_dimensions, prompt_upsampling],
+        inputs=[prompt, input_images, seed, randomize_seed, width, height, num_inference_steps, guidance_scale, prompt_upsampling],
         outputs=[result, seed]
     )