fixed appearance and added more images

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

app.py

@@ -6,12 +6,12 @@ import torch.nn.functional as F
 import os
 import glob
 
-from archs import create_model, resume_model
+from archs import create_model, load_model
 
 # -------- Detect folders & images (assets/<folder>) --------
 IMG_EXTS = (".png", ".jpg", ".jpeg", ".bmp", ".webp")
 
-def list_subfolders(base="assets"):
+def list_subfolders(base="examples"):
     """Return a sorted list of immediate subfolders inside base."""
     if not os.path.isdir(base):
         return []
@@ -19,14 +19,15 @@ def list_subfolders(base="assets"):
     return subs
 
 def list_images(folder):
-    """Return full paths of images inside assets/<folder>."""
-    paths = sorted(glob.glob(os.path.join("assets", folder, "*")))
+    """Return full paths of images inside examples/<folder>."""
+    paths = sorted(glob.glob(os.path.join("examples", folder, "*")))
     return [p for p in paths if p.lower().endswith(IMG_EXTS)]
 
 # -------- Folder/Gallery interactions --------
 def update_gallery(folder):
     """Given a folder name, return the gallery items (list of image paths) and store the same list in state."""
     files = list_images(folder)
+    print(files)
     return gr.update(value=files, visible=True), files 
 
 def load_from_gallery(evt: gr.SelectData, current_files):
@@ -35,6 +36,7 @@ def load_from_gallery(evt: gr.SelectData, current_files):
     if not current_files or idx is None or idx >= len(current_files):
         return gr.update()
     path = current_files[idx]
+    print(path)
     return Image.open(path)
 
 
@@ -59,7 +61,9 @@ tensor_to_pil = transforms.ToPILImage()
 model = create_model(model_opt, device)
 
 checkpoints = torch.load(PATH_MODEL, map_location=device, weights_only=False)
-model = resume_model(model, PATH_MODEL, device)
+model = load_model(model, PATH_MODEL, device)
+
+model.eval()
 
 def pad_tensor(tensor, multiple = 16):
     '''pad the tensor to be multiple of some number'''
@@ -88,12 +92,14 @@ def process_img(image, task_label = 'auto'):
     task_label = LABEL_TO_TASK.get(task_label, 'auto')
     tensor = pil_to_tensor(image).unsqueeze(0).to(device)
     _, _, H, W = tensor.shape
-    
+    print('Using task:', task_label)
     tensor = pad_tensor(tensor)
 
     with torch.no_grad():
-        output = model(tensor, task_label)
+        output_dict = model(tensor, task_label)
 
+    output = output_dict['output']
+    # print(output.shape)
     output = torch.clamp(output, 0., 1.)
     output = output[:,:, :H, :W].squeeze(0)    
     return tensor_to_pil(output)
@@ -117,17 +123,12 @@ Available code at [github](https://github.com/cidautai/DeMoE). More information
 <br>
 '''
 
-# examples = [['examples/1POA1811.png'],
-#             ['examples/12_blur.png'],
-#             ['examples/0031.png'],
-#             ['examples/000143.png'],
-#             ['examples/blur_4.png']]
-
 css = """
-    .image-frame img, .image-container img {
-        width: auto;
-        height: auto;
-        max-width: none;
+.fitbox img,
+    .fitbox canvas {
+    width: 100% !important;
+    height: 100% !important;
+    object-fit: contain !important;
     }
 """
 
@@ -150,6 +151,7 @@ examples_synth_global_motion = list_basenames("synth_global_motion")
 examples_local_motion      = list_basenames("local_motion")
 examples_defocus    = list_basenames("defocus")
 
+# print(examples_defocus, examples_global_motion, examples_low_light, examples_synth_global_motion, examples_local_motion)
 # -----------------------------
 # Gradio Blocks layout
 # -----------------------------
@@ -158,12 +160,12 @@ with gr.Blocks(css=css, title=title) as demo:
 
     with gr.Row():
         # Input image and the task selector (Radio)
-        inp_img = gr.Image(type='pil', label='input')
+        inp_img = gr.Image(type='pil', label='input', height=320)
         # Output image and action button
-        out_img = gr.Image(type='pil', label='output')
+        out_img = gr.Image(type='pil', label='output', height=320)
     task_selector = gr.Radio(
     choices=TASK_LABELS,
-    value="auto",
+    value="Auto",
     label="Blur type"
 )
 
@@ -181,6 +183,7 @@ with gr.Blocks(css=css, title=title) as demo:
     with gr.Row():
         # List folders found in ./assets
         folders = list_subfolders("examples")
+        print(folders)
         folder_radio = gr.Radio(choices=folders, label="Examples Folders", interactive=True)
 
         gallery = gr.Gallery(

archs/__init__.py

@@ -25,34 +25,70 @@ def create_model(opt, device):
     
     return model
 
-def load_weights(model, model_weights):
-    '''
-    Loads the weights of a pretrained model, picking only the weights that are
-    in the new model.
-    '''
-    new_weights = model.state_dict()
-    new_weights.update({k: v for k, v in model_weights.items() if k in new_weights})
+# def load_weights(model, model_weights):
+#     '''
+#     Loads the weights of a pretrained model, picking only the weights that are
+#     in the new model.
+#     '''
+#     new_weights = model.state_dict()
+#     new_weights.update({k: v for k, v in model_weights.items() if k in new_weights})
     
-    model.load_state_dict(new_weights)
+#     new_weights = {key.replace('module.', ''): value for key, value in new_weights.items()}
+#     print(new_weights.keys())
+#     print(model.state_dict().keys())
+#     model.load_state_dict(new_weights, strict= True)
+
+
+#     total_checkpoint_keys = len(model_weights)
+#     total_model_keys = len(new_weights)
+#     matching_keys = len(set(model_weights.keys()) & set(new_weights.keys()))
+
+#     print(f"Total keys in checkpoint: {total_checkpoint_keys}")
+#     print(f"Total keys in model state dict: {total_model_keys}")
+#     print(f"Number of matching keys: {matching_keys}")
+
 
+#     return model
+
+def strip_prefixes(sd: dict, prefixes=("module.", "model.", "ema.", "net.", "netG.", "generator.")) -> dict:
+    out = {}
+    for k, v in sd.items():
+        nk = k
+        for p in prefixes:
+            if nk.startswith(p):
+                nk = nk[len(p):]
+                break
+        out[nk] = v
+    return out
+
+# ===== quita DDP y local_rank, usa un device único =====
+def load_model(model, path_weights: str, device: torch.device):
+    # siempre carga en CPU y luego mueve
+    ckpt = torch.load(path_weights, map_location='cpu', weights_only=False)
+    # intenta varias claves habituales; si no, usa el dict tal cual
+    sd = ckpt.get("params") or ckpt.get("model_state_dict") or ckpt.get("state_dict") or ckpt
+    sd = strip_prefixes(sd)
+    missing, unexpected = model.load_state_dict(sd, strict=False)
+    print(f"[DeMoE] load_state: missing={len(missing)}, unexpected={len(unexpected)}")
+    model = model.to(device=device, dtype=torch.float32).eval()
     return model
 
-def resume_model(model,
-                 path_model, 
-                 device):
+# def resume_model(model,
+#                  path_model, 
+#                  device):
     
-    '''
-    Returns the loaded weights of model and optimizer if resume flag is True
-    '''
+#     '''
+#     Returns the loaded weights of model and optimizer if resume flag is True
+#     '''
 
-    checkpoints = torch.load(path_model, map_location=device, weights_only=False)
-    weights = checkpoints['params']
-    model = load_weights(model, model_weights=weights)
+#     checkpoints = torch.load(path_model, map_location=device, weights_only=False)
+#     weights = checkpoints['params']
+#     model = load_weights(model, model_weights=weights)
 
-    return model
+#     return model
 
 
-__all__ = ['create_model', 'resume_model', 'load_weights']
+__all__ = ['create_model', 'load_model']