fixed the randomness issue

app.py

@@ -1,9 +1,12 @@
 import gradio as gr
 import os, requests
+import numpy as np
 from inference import setup_model, colorize_grayscale, predict_anchors
 
-os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/disco-beta.pth.rar")
-os.rename("disco-beta.pth.rar", "./checkpoints/disco-beta.pth.rar")
+RUN_MODE = "remote"
+if RUN_MODE != "local":
+    os.system("wget https://huggingface.co/menghanxia/disco/resolve/main/disco-beta.pth.rar")
+    os.rename("disco-beta.pth.rar", "./checkpoints/disco-beta.pth.rar")
 
 ## step 1: set up model
 device = "cpu"
@@ -30,7 +33,7 @@ def switch_states(is_checked):
 demo = gr.Blocks(title="DISCO")
 with demo:
     gr.Markdown(value="""
-                    **Gradio demo for DISCO: Disentangled Image Colorization via Global Anchors. [Project Page](https://menghanxia.github.io/projects/disco.html)**.
+                    **Gradio demo for DISCO: Disentangled Image Colorization via Global Anchors**. Check our project page [*Here*](https://menghanxia.github.io/projects/disco.html).
                     """)
     with gr.Row():
         with gr.Column(scale=1):
@@ -39,13 +42,13 @@ with demo:
             with gr.Row():
                 Num_anchor = gr.Number(type="int", value=8, label="Num. of anchors (3~14)")
                 Radio_resolution = gr.Radio(type="index", choices=["Low (256x256)", "High (512x512)"], \
-                                            label="Colorization resolution", value="Low (256x256)")
+                                            label="Colorization resolution (Low is more stable)", value="Low (256x256)")
             Ckeckbox_editable = gr.Checkbox(default=False, label='Show editable anchors')
             with gr.Row():
                 Button_show_anchor = gr.Button(value="Predict anchors", visible=False)
                 Button_run = gr.Button(value="Colorize")
         with gr.Column(scale=1):
-            Image_output = gr.Image(type="numpy", label="Output", shape=[100,100])
+            Image_output = gr.Image(type="numpy", label="Output").style(height=480)
     
     Ckeckbox_editable.change(fn=switch_states, inputs=Ckeckbox_editable, outputs=[Image_anchor, Button_show_anchor])
     Button_show_anchor.click(fn=click_predanchors, inputs=[Image_input, Num_anchor, Radio_resolution, Ckeckbox_editable], outputs=Image_anchor)
@@ -55,14 +58,16 @@ with demo:
     gr.Markdown(value="""    
                     **Guideline**
                     1. upload your image;
-                    2. Set up the arguments: "Num. of anchors" and "Colorization resolution";
-                    3. Run the colorization (two modes supported):
-                        - **Automatic**: click "Colorize" to get the automatically colorized output.
-                        - **Editable**: check ""Show editable anchors" and click "Predict anchors". Then, modify the colors of the predicted anchors (only anchor region will be used). Finally, click "Colorize" to get the result.
+                    2. set up the arguments: "Num. of anchors" and "Colorization resolution";
+                    3. run the colorization (two modes supported):
+                        - *Automatic mode*: click "Colorize" to get the automatically colorized output.
+                        - *Editable mode*: check ""Show editable anchors" and click "Predict anchors". Then, modify the colors of the predicted anchors (only anchor region will be used). Finally, click "Colorize" to get the result.
                     """)
     gr.HTML(value="""
-                <p style='text-align: center'><a href='https://menghanxia.github.io/projects/disco.html' target='_blank'>DISCO Project Page</a> | <a href='https://github.com/MenghanXia/DisentangledColorization' target='_blank'>Github Repo</a></p>
+                <p style="text-align:center; color:orange"><a href='https://menghanxia.github.io/projects/disco.html' target='_blank'>DISCO Project Page</a> | <a href='https://github.com/MenghanXia/DisentangledColorization' target='_blank'>Github Repo</a></p>
                     """)
-    
-#demo.launch(server_name='9.134.253.83',server_port=7788)
-demo.launch()
+
+if RUN_MODE == "local":
+    demo.launch(server_name='9.134.253.83',server_port=7788)
+else:
+    demo.launch()

inference.py

@@ -11,17 +11,13 @@ from utils import util
 
 
 def setup_model(checkpt_path, device="cuda"):
-    seed = 130
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed(seed)
     #print('--------------', torch.cuda.is_available())
     """Load the model into memory to make running multiple predictions efficient"""
     colorLabeler = basic.ColorLabel(device=device)
     colorizer = model.AnchorColorProb(inChannel=1, outChannel=313, enhanced=True, colorLabeler=colorLabeler)
     colorizer = colorizer.to(device)
     #checkpt_path = "./checkpoints/disco-beta.pth.rar"
-    assert os.path.exists(checkpt_path)
+    assert os.path.exists(checkpt_path), "No checkpoint found!"
     data_dict = torch.load(checkpt_path, map_location=torch.device('cpu'))
     colorizer.load_state_dict(data_dict['state_dict'])
     colorizer.eval()
@@ -89,8 +85,8 @@ def predict_anchors(colorizer, color_class, rgb_img, n_anchors, is_high_res, is_
     n_anchors = min(n_anchors, 14)
     target_res = (512,512) if is_high_res else (256,256)
     input_grays, input_colors, org_grays = prepare_data(rgb_img, target_res)
-    input_grays = input_grays.cuda(non_blocking=True)
-    input_colors = input_colors.cuda(non_blocking=True)
+    input_grays = input_grays.to(device)
+    input_colors = input_colors.to(device)
                 
     sampled_T, sp_size = 0, 16
     pal_logit, ref_logit, enhanced_ab, affinity_map, spix_colors, hint_mask = colorizer(input_grays, \

models/clusterkit.py

@@ -103,6 +103,7 @@ def initialize(X, num_clusters):
     :param num_clusters: (int) number of clusters
     :return: (np.array) initial state
     """
+    np.random.seed(1)
     num_samples = len(X)
     indices = np.random.choice(num_samples, num_clusters, replace=False)
     initial_state = X[indices]