import flair fix

app.py

@@ -14,8 +14,8 @@ project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "."))
 if project_root not in sys.path:
     sys.path.insert(0, project_root)
 
-from flair.pipelines import model_loader
-from flair import var_post_samp, degradations
+from src.flair.pipelines import model_loader
+from src.flair import var_post_samp, degradations
 
 
 
@@ -283,7 +283,7 @@ def inpaint_image(image_editor_output, prompt_text, fixed_seed_value, use_random
 
         current_seed = None
         if use_random_seed:
-            current_seed = None # load_config_for_inference will generate a random seed
+            current_seed = random.randint(0, 2**32 - 1)
         else:
             try:
                 current_seed = int(fixed_seed_value)
@@ -514,15 +514,16 @@ if os.path.exists(demo_images_dir):
                         metadata = json.load(f)
                     task = metadata.get("task_type")
                     prompt = metadata.get("prompt", "")
+                    n_steps = metadata.get("num_steps", 50)
                     if task == "Super Resolution":
-                        example_list_sr.append([image_path, prompt, task])
+                        example_list_sr.append([image_path, prompt, task, n_steps])
                     else:
                         image_editor_input = {
                             "background": image_path,
                             "layers": [mask_path],
                             "composite": None  # Add this key to satisfy ImageEditor's as_example processing
                         }
-                        example_list_inp.append([image_editor_input, prompt, task])
+                        example_list_inp.append([image_editor_input, prompt, task, n_steps])
                     
                     # Structure for ImageEditor: { "background": filepath, "layers": [filepath], "composite": None }
                     
@@ -553,10 +554,14 @@ if __name__ == "__main__":
         sys.exit(1)
 
     # --- Define Gradio UI using gr.Blocks after globals are initialized ---
-    title_str = "Solving Inverse Problems with FLAIR: Inpainting Demo"
+    title_str = "Solving Inverse Problems with FLAIR"
     description_str = """
 Select a task (Inpainting or Super Resolution) and upload an image.
-For Inpainting, draw a mask on the image to specify the area to be filled. We observed that our model can event solve simple editing task, if provided with an appropriate prompt. For large masks the step size might need to be adjusted to e.g. 80.
+
+For Inpainting, draw a mask on the image to specify the area to be filled. 
+
+We observed that our model can event solve simple editing task, if provided with an appropriate prompt. For large masks the step size might need to be adjusted to e.g. 80.
+
 For Super Resolution, upload a low-resolution image and select the upscaling factor. Images are always upscaled to 768x768 pixels. Therefore, for x12 superresolution, the input image must be 64x64 pixels. You can also upload a high resolution image which will be downscaled to the correct input size. 
 Use the slider to compare the low resolution input image with the super-resolved output.
 
@@ -729,43 +734,55 @@ Use the slider to compare the low resolution input image with the super-resolved
 
         gr.Markdown("---") # Separator
         gr.Markdown("### Click an example to load:")
-        def load_example(example_data, prompt, task):
-            image_editor_input = example_data[0]
-            prompt_value = example_data[1]
-            if task == "Inpainting":
-                image_editor.clear()  # Clear current image and mask
-                if image_editor_input and image_editor_input.get("background"):
-                    image_editor.upload_image(image_editor_input["background"])
-                if image_editor_input and image_editor_input.get("layers"):
-                    for layer in image_editor_input["layers"]:
-                        image_editor.upload_mask(layer)
-            elif task == "Super Resolution":
-                image_input.clear()
-                image_input.upload_image(image_editor_input)
-                
-            # Set the prompt
-            prompt_text.value = prompt_value
-            # Optionally, set a random seed and guidance scale
-            seed_slider.value = random.randint(0, 2**32 - 1)
-            guidance_scale_slider.value = default_guidance_scale
-            # Set the task selector from the example
-            task_selector.set_value(task)
-            update_visibility(task)  # Update visibility based on task
 
-        with gr.Row():
-            gr.Examples(
-                examples=example_list_sr,
-                inputs=[image_input, prompt_text, task_selector],
-                label="Super Resolution Examples",
-                fn=load_example,
-            )
-        with gr.Row():
-            gr.Examples(
-                examples=example_list_inp,
-                inputs=[image_editor, prompt_text, task_selector],
-                label="Inpainting Examples",
-                fn=load_example,
-            )
+        # --- GALLERY FOR SUPER RESOLUTION EXAMPLES ---
+        sr_gallery_items = [[ex[0], f"Prompt: {ex[1]} Steps: {ex[3]}"] for ex in example_list_sr]
+        sr_gallery = gr.Gallery(
+            value=sr_gallery_items,
+            label="Super Resolution Examples",
+            columns=4,
+            height="auto",
+            visible=True
+        )
+
+        # --- GALLERY FOR INPAINTING EXAMPLES ---
+        inp_gallery_items = [[ex[0]["background"], f"Prompt: {ex[1]} Steps: {ex[3]}"] for ex in example_list_inp]
+        inp_gallery = gr.Gallery(
+            value=inp_gallery_items,
+            label="Inpainting Examples",
+            columns=4,
+            height="auto",
+            visible=True
+        )
+
+        def on_sr_gallery_select(evt: gr.SelectData):
+            idx = evt.index
+            ex = example_list_sr[idx]
+            image_input.value = ex[0]
+            prompt_text.value = ex[1]
+            task_selector.value = ex[2]
+            num_steps_slider.value = ex[3]
+            update_visibility(ex[2])
+            return [image_input, prompt_text, task_selector, num_steps_slider]
+
+        def on_inp_gallery_select(evt: gr.SelectData):
+            idx = evt.index
+            ex = example_list_inp[idx]
+            image_editor.value = ex[0]
+            prompt_text.value = ex[1]
+            task_selector.value = ex[2]
+            num_steps_slider.value = ex[3]
+            update_visibility(ex[2])
+            return [image_editor, prompt_text, task_selector, num_steps_slider]
+
+        sr_gallery.select(
+            fn=on_sr_gallery_select,
+            outputs=[image_input, prompt_text, task_selector, num_steps_slider]
+        )
+        inp_gallery.select(
+            fn=on_inp_gallery_select,
+            outputs=[image_editor, prompt_text, task_selector, num_steps_slider]
+        )
 
     # --- End of Gradio UI definition ---
 

demo_images/demo_2_meta.json

@@ -2,6 +2,6 @@
     "prompt": "a high quality image of a face.",
     "seed_on_slider": 3211750901,
     "use_random_seed_checkbox": false,
-    "num_steps": 50,
+    "num_steps": 75,
     "task_type": "Inpainting"
 }

requirements.txt

@@ -16,4 +16,3 @@ sentencepiece
 protobuf
 accelerate
 gradio
-gradio_imageslider

src/flair/degradations.py

@@ -2,7 +2,7 @@ import torch
 import numpy as np
 from munch import munchify
 from scipy.ndimage import distance_transform_edt
-from flair.functions.degradation import get_degradation
+from src.flair.functions.degradation import get_degradation
 import torchvision 
 
 class BaseDegradation(torch.nn.Module):

src/flair/functions/degradation.py

@@ -2,9 +2,9 @@ import numpy as np
 import torch
 from munch import Munch
 
-from flair.functions import svd_operators as svd_op
-from flair.functions import measurements
-from flair.utils.inpaint_util import MaskGenerator
+from src.flair.functions import svd_operators as svd_op
+from src.flair.functions import measurements
+from src.flair.utils.inpaint_util import MaskGenerator
 
 __DEGRADATION__ = {}
 
@@ -187,7 +187,7 @@ def deg_deblur_guass_general(deg_config, device):
     return A_funcs
 
 
-from flair.functions.jpeg import jpeg_encode, jpeg_decode
+from src.flair.functions.jpeg import jpeg_encode, jpeg_decode
 
 class JPEGOperator():
     def __init__(self, qf: int, device):

src/flair/functions/measurements.py

@@ -6,15 +6,15 @@ from functools import partial
 from torch.nn import functional as F
 from torchvision import torch
 
-from flair.utils.blur_util import Blurkernel
-from flair.utils.img_util import fft2d
+from src.flair.utils.blur_util import Blurkernel
+from src.flair.utils.img_util import fft2d
 import numpy as np
-from flair.utils.resizer import Resizer
-from flair.utils.utils_sisr import pre_calculate_FK, pre_calculate_nonuniform
+from src.flair.utils.resizer import Resizer
+from src.flair.utils.utils_sisr import pre_calculate_FK, pre_calculate_nonuniform
 from torch.fft import fft2, ifft2
 
 
-from flair.motionblur.motionblur import Kernel
+from src.flair.motionblur.motionblur import Kernel
 
 # =================
 # Operation classes

src/flair/pipelines/model_loader.py

@@ -10,7 +10,7 @@ from diffusers import (
 from diffusers import AutoencoderTiny
 
 
-from flair.pipelines import sd3
+from src.flair.pipelines import sd3
 
 
 

src/flair/pipelines/sd3.py

@@ -1,7 +1,7 @@
 import torch
 from typing import Dict, Any
 from diffusers.pipelines.stable_diffusion_3 import pipeline_stable_diffusion_3
-from flair.pipelines import utils
+from src.flair.pipelines import utils
 import tqdm
 
 class SD3Wrapper(pipeline_stable_diffusion_3.StableDiffusion3Pipeline):

src/flair/utils/blur_util.py

@@ -3,7 +3,7 @@ from torch import nn
 import numpy as np
 import scipy
 
-from flair.utils.motionblur import Kernel as MotionKernel
+from src.flair.utils.motionblur import Kernel as MotionKernel
 
 class Blurkernel(nn.Module):
     def __init__(self, blur_type='gaussian', kernel_size=31, std=3.0, device=None):

src/flair/var_post_samp.py

@@ -3,7 +3,7 @@ import numpy as np
 import sys
 import os
 import tqdm
-from flair import degradations
+from src.flair import degradations
 import torchvision
 
 def total_variation_loss(x):