added progress bar

app.py

@@ -10,8 +10,6 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 device = "mps" if torch.backends.mps.is_available() else device
 
 image_size = 128
-upscale = False
-clicked = False
 
 
 transform = transforms.Compose(
@@ -23,49 +21,26 @@ transform = transforms.Compose(
 )
 
 
-def make_scribbles(sketch, scribbles):
-    # get the value that occurs most often in the scribbles
-    sketch = transforms.Resize((image_size, image_size))(sketch)
-    scribbles = transforms.Resize((image_size, image_size))(scribbles)
-
-    grey_tensor = torch.tensor(0.49803922, device=device)
-
-    grey_tensor = grey_tensor.expand(3, image_size, image_size)
-
-    sketch = transforms.ToTensor()(sketch).to(device)
-    scribbles = transforms.ToTensor()(scribbles).to(device)
-
-    scribble_where_grey_mask = torch.eq(scribbles, grey_tensor)
-
-    merged = torch.where(scribble_where_grey_mask, sketch, scribbles)
-
-    return transforms.Lambda(lambda t: (t * 2) - 1)(sketch), transforms.Lambda(
-        lambda t: (t * 2) - 1
-    )(merged)
-
-
-def process_images(sketch, scribbles, sampling_steps, is_scribbles, seed_nr, upscale):
-    global clicked
-    clicked = True
+def process_images(
+    sketch,
+    scribbles,
+    sampling_steps,
+    seed_nr,
+    upscale,
+    progress=gr.Progress(),
+):
     w, h = sketch.size
 
-    if is_scribbles:
-        sketch, scribbles = make_scribbles(sketch, scribbles)
-
-    else:
-        sketch = transform(sketch.convert("RGB"))
-        scribbles = transform(scribbles.convert("RGB"))
+    sketch = transform(sketch.convert("RGB"))
+    scribbles = transform(scribbles.convert("RGB"))
 
     if upscale:
-        output = transforms.Resize((h, w))(
-            sample(sketch, scribbles, sampling_steps, seed_nr)
+        return transforms.Resize((h, w))(
+            sample(sketch, scribbles, sampling_steps, seed_nr, progress)
         )
-        clicked = False
-        return output
+
     else:
-        output = sample(sketch, scribbles, sampling_steps, seed_nr)
-        clicked = False
-        return output
+        return sample(sketch, scribbles, sampling_steps, seed_nr, progress)
 
 
 theme = gr.themes.Monochrome()
@@ -87,7 +62,7 @@ with gr.Blocks(theme=theme) as demo:
                 "By default the scribbles are assumed to be merged with the sketch, if they appear on a grey background check the box below. "
                 "</p>"
             )
-            is_scribbles = gr.Checkbox(label="Is Scribbles", value=False)
+
         with gr.Column():
             output = gr.Image(type="pil", label="Output")
             upscale_info = gr.Markdown(
@@ -96,14 +71,12 @@ with gr.Blocks(theme=theme) as demo:
                 "</p>"
             )
             upscale_button = gr.Checkbox(label="Stretch", value=False)
+
     with gr.Row():
         with gr.Column():
             seed_slider = gr.Number(
-                label="Random Seed 🎲",
-                value=random.randint(
-                    1,
-                    1000,
-                ),
+                label="Random Seed 🎲 (if the image generated is not to your liking, simply use another seed)",
+                value=random.randint(0, 10000),
             )
 
         with gr.Column():
@@ -111,12 +84,12 @@ with gr.Blocks(theme=theme) as demo:
                 minimum=1,
                 maximum=250,
                 step=1,
-                label="DDPM Sampling Steps 🔄",
+                label="DDPM Sampling Steps 🔄 (the higher the number of steps the higher the quality of the images)",
                 value=50,
             )
 
     with gr.Row():
-        generate_button = gr.Button(value="Generate", interactive=not clicked)
+        generate_button = gr.Button(value="Generate")
     with gr.Row():
         generate_info = gr.Markdown(
             "<p style='text-align: center; font-size: 16px;'>"
@@ -130,13 +103,13 @@ with gr.Blocks(theme=theme) as demo:
             sketch_input,
             scribbles_input,
             sampling_slider,
-            is_scribbles,
             seed_slider,
             upscale_button,
         ],
         outputs=output,
-        show_progress=True,
+        concurrency_limit=1,
+        trigger_mode="once",
     )
 
 if __name__ == "__main__":
-    demo.launch(max_threads=1)
+    demo.queue().launch(max_threads=1)

load_model.py

@@ -8,6 +8,7 @@ from torchvision import transforms
 import pathlib
 from torchvision.utils import save_image
 from safetensors.torch import load_model, save_model
+import time as tm
 
 
 denoising_timesteps = 4000
@@ -61,7 +62,7 @@ else:
     raise Exception("No model files found in the folder.")
 
 
-def sample(sketch, scribbles, sampling_steps, seed_nr):
+def sample(sketch, scribbles, sampling_steps, seed_nr, progress):
     torch.manual_seed(seed_nr)
 
     noise_scheduler = DDPMScheduler(
@@ -80,9 +81,9 @@ def sample(sketch, scribbles, sampling_steps, seed_nr):
 
         noise_for_plain = torch.randn_like(sketch, device=device)
 
-        for i, t in tqdm(
-            enumerate(noise_scheduler.timesteps),
-            total=len(noise_scheduler.timesteps),
+        for t in progress.tqdm(
+            noise_scheduler.timesteps,
+            desc="Sampling",
         ):
             noise_for_plain = noise_scheduler.scale_model_input(noise_for_plain, t).to(
                 device
@@ -105,6 +106,8 @@ def sample(sketch, scribbles, sampling_steps, seed_nr):
                 noise_for_plain,
             ).prev_sample
 
+            tm.sleep(0.01)
+
     sample = torch.clamp((noise_for_plain / 2) + 0.5, 0, 1)
 
     return transforms.ToPILImage()(sample[0].cpu())