fix: less verbose on gif and progress output (#127)

* fix: less verbose on gif and progress output

* style: fix overload and cli autocomplete

Co-authored-by: Jina Dev Bot <dev-bot@jina.ai>

Author: hanxiao

discoart/config.py

@@ -27,7 +27,7 @@
 ) as ymlfile:
     cut_schedules = yaml.load(ymlfile, Loader=Loader)
 
-_legacy_args = {'clip_sequential_evaluation', 'fuzzy_prompt'}
+_legacy_args = {'clip_sequential_evaluation', 'fuzzy_prompt', 'display_rate'}
 
 
 def load_config(

discoart/create.py

@@ -39,7 +39,6 @@ def create(
     diffusion_model: Optional[str] = '512x512_diffusion_uncond_finetune_008100',
     diffusion_model_config: Optional[Dict[str, Any]] = None,
     diffusion_sampling_mode: Optional[str] = 'ddim',
-    display_rate: Optional[int] = None,
     eta: Optional[float] = 0.8,
     gif_fps: Optional[int] = 20,
     gif_size_ratio: Optional[float] = 0.5,
@@ -113,9 +112,8 @@ def create(**kwargs) -> Optional['DocumentArray']:
     :param diffusion_model: Diffusion_model of choice. Note that you don't have to write the full name of the diffusion model, e.g. any prefix is enough.To use a listed all diffusion models, you can do:```pythonfrom discoart import createcreate(diffusion_model='portrait_generator', ...)```
     :param diffusion_model_config: [DiscoArt] The customized diffusion model config as a dictionary, if specified will override the values with the same name in the default model config.
     :param diffusion_sampling_mode: Two alternate diffusion denoising algorithms. ddim has been around longer, and is more established and tested.  plms is a newly added alternate method that promises good diffusion results in fewer steps, but has not been as fully tested and may have side effects. This new plms mode is actively being researched in the #settings-and-techniques channel in the DD Discord.
-    :param display_rate: [DiscoArt] Display rate is deprecated in DiscoArt as it is always 1, meaning display is always in real-time. There is no need to worry on the speed as the rendering happens in another thread. To control the save rate, use the `save_rate` parameter.Setting this will override the `save_rate` parameter.
     :param eta: eta (greek letter η) is a diffusion model variable that mixes in a random amount of scaled noise into each timestep. 0 is no noise, 1.0 is more noise. As with most DD parameters, you can go below zero for eta, but it may give you unpredictable results. The steps parameter has a close relationship with the eta parameter. If you set eta to 0, then you can get decent output with only 50-75 steps. Setting eta to 1.0 favors higher step counts, ideally around 250 and up. eta has a subtle, unpredictable effect on image, so you’ll need to experiment to see how this affects your projects.
-    :param gif_fps: [DiscoArt] The frame rate of the generated GIF.
+    :param gif_fps: [DiscoArt] The frame rate of the generated GIF. Set it to -1 for not saving GIF.
     :param gif_size_ratio: [DiscoArt] The relative size vs. the original image, small size ratio gives smaller file size.
     :param init_document: [DiscoArt] Use a Document object as the initial state for DD: its ``.tags`` will be used as parameters, ``.uri`` (if present) will be used as init image.
     :param init_image: Recall that in the image sequence above, the first image shown is just noise.  If an init_image is provided, diffusion will replace the noise with the init_image as its starting state.  To use an init_image, upload the image to the Colab instance or your Google Drive, and enter the full image path here. If using an init_image, you may need to increase skip_steps to ~ 50% of total steps to retain the character of the init. See skip_steps above for further discussion.
@@ -129,7 +127,7 @@ def create(**kwargs) -> Optional['DocumentArray']:
     :param randomize_class: Controls whether the imagenet class is randomly changed each iteration
     :param range_scale: Optional, set to zero to turn off.  Used for adjustment of color contrast.  Lower range_scale will increase contrast. Very low numbers create a reduced color palette, resulting in more vibrant or poster-like images. Higher range_scale will reduce contrast, for more muted images.[DiscoArt] Can be scheduled via syntax `[val1]*400+[val2]*600`.
     :param sat_scale: Saturation scale. Optional, set to zero to turn off.  If used, sat_scale will help mitigate oversaturation. If your image is too saturated, increase sat_scale to reduce the saturation.[DiscoArt] Can be scheduled via syntax `[val1]*400+[val2]*600`.
-    :param save_rate: [DiscoArt] The number of steps to save intermediate results. It is a replacement to original `display_rate` parameter.
+    :param save_rate: [DiscoArt] The number of steps to save intermediate results. It is a replacement to original `display_rate` parameter. Set it to -1 for not saving any intermediate result.
     :param seed: Deep in the diffusion code, there is a random number ‘seed’ which is used as the basis for determining the initial state of the diffusion.  By default, this is random, but you can also specify your own seed.  This is useful if you like a particular result and would like to run more iterations that will be similar. After each run, the actual seed value used will be reported in the parameters report, and can be reused if desired by entering seed # here.  If a specific numerical seed is used repeatedly, the resulting images will be quite similar but not identical.
     :param skip_augs: Controls whether to skip torchvision augmentations.[DiscoArt] Can be scheduled via syntax `[val1]*400+[val2]*600`.
     :param skip_event: [DiscoArt] A multiprocessing/asyncio/threading.Event that once set, will skip the current run and move to the next run as defined in `n_batches`.

discoart/persist.py

@@ -29,6 +29,7 @@ def _sample(
     output_dir,
     is_sampling_done,
     is_save_step,
+    is_save_gif,
 ):
     with threading.Lock():
         is_sampling_done.clear()
@@ -57,11 +58,13 @@ def _sample(
                     )
 
                 da[k].chunks.append(c)
-                # root doc always update with the latest progress
-                da[k].uri = c.uri
-            else:
+
+            if is_save_gif:
                 da_gif[k].chunks.append(c)
 
+            # root doc always update with the latest progress
+            da[k].uri = c.uri
+
             da[k].tags['_status'] = {
                 'completed': cur_t == -1,
                 'cur_t': cur_t,

discoart/resources/default.yml

@@ -36,7 +36,6 @@ cut_innercut: "[4]*400+[12]*600"
 cut_icgray_p: "[0.2]*400+[0]*600"
 cut_ic_pow: 1.
 
-display_rate:
 save_rate: 20
 gif_fps: 20
 gif_size_ratio: 0.5

discoart/resources/docstrings.yml

@@ -121,13 +121,8 @@ cut_innercut: |
 cut_icgray_p: |
   In addition to the overall cut schedule, a portion of the cuts can be set to be grayscale instead of color. This may help with improved definition of shapes and edges, especially in the early diffusion steps where the image structure is being defined.  
 
-display_rate: |
-  [DiscoArt] Display rate is deprecated in DiscoArt as it is always 1, meaning display is always in real-time. 
-  
-  There is no need to worry on the speed as the rendering happens in another thread. To control the save rate, use the `save_rate` parameter.
-  Setting this will override the `save_rate` parameter.
 save_rate: |
-  [DiscoArt] The number of steps to save intermediate results. It is a replacement to original `display_rate` parameter.
+  [DiscoArt] The number of steps to save intermediate results. It is a replacement to original `display_rate` parameter. Set it to -1 for not saving any intermediate result.
 
 n_batches: |
   This variable sets the number of still images you want DD to create.  If you are using an animation mode (see below for details) DD will ignore n_batches and create a single set of animated frames based on the animation settings.
@@ -220,7 +215,7 @@ text_clip_on_cpu: |
   [DiscoArt] Place text transformers of CLIP models on CPU. This saves more VRAM and will not hurt the speed at all on T4, P100, 3090; however, there are few community members report issue on V100 when it is `False`.
 
 gif_fps: |
-  [DiscoArt] The frame rate of the generated GIF.
+  [DiscoArt] The frame rate of the generated GIF. Set it to -1 for not saving GIF.
 
 gif_size_ratio: |
   [DiscoArt] The relative size vs. the original image, small size ratio gives smaller file size.

discoart/runner.py

@@ -1,6 +1,7 @@
 import copy
 import os.path
 import random
+import tempfile
 import threading
 
 import clip
@@ -301,9 +302,9 @@ def cond_fn(x, t, **kwargs):
             x_is_NaN = True
             grad = torch.zeros_like(x)
             logger.warning(
-                f'NaN detected in grad at the diffusion inner-step {num_step}, '
-                f'if this message continues to show up, '
-                f'then your image is not updated and further steps are unnecessary.'
+                f'NaN detected in grad at the diffusion inner-step {num_step}, no panic. '
+                f'However, if this message continues to show up *in a row*, '
+                f'then your generation is ill-conditioned and image will not updated, further steps are unnecessary.'
             )
 
         r_grad = grad
@@ -379,7 +380,6 @@ def cond_fn(x, t, **kwargs):
                 _handlers,
                 _redraw_fn,
                 args,
-                output_dir,
                 _nb,
             )
         free_memory()
@@ -443,9 +443,9 @@ def cond_fn(x, t, **kwargs):
 
                 cur_t -= 1
 
-                is_save_step = (
-                    j % (args.display_rate or args.save_rate) == 0 or cur_t == -1
-                )
+                is_save_step = args.save_rate > 0 and j % args.save_rate == 0
+                is_complete = cur_t == -1
+
                 threads.append(
                     _sample_thread(
                         sample,
@@ -459,6 +459,7 @@ def cond_fn(x, t, **kwargs):
                         output_dir,
                         is_busy_evs[0],
                         is_save_step,
+                        args.gif_fps > 0,
                     )
                 )
 
@@ -473,13 +474,15 @@ def cond_fn(x, t, **kwargs):
                             args.gif_size_ratio,
                         )
                     )
+
+                if is_complete or is_save_step:
                     threads.extend(
                         _persist_thread(
                             da_batches,
                             args.name_docarray,
                             is_busy_evs[1:],
                             is_busy_evs[0],
-                            is_completed=cur_t == -1,
+                            is_completed=is_complete,
                         )
                     )
 
@@ -497,19 +500,21 @@ def cond_fn(x, t, **kwargs):
     return da_batches
 
 
-def redraw_widget(_handlers, _redraw_fn, args, output_dir, _nb):
+def redraw_widget(_handlers, _redraw_fn, args, _nb):
     _handlers.progress.max = args.n_batches
     _handlers.progress.value = _nb + 1
     _handlers.progress.description = f'Baking {_nb + 1}/{args.n_batches}: '
 
-    svg0 = os.path.join(output_dir, 'config.svg')
-    save_config_svg(args, svg0, only_non_default=True)
-    d = Document(uri=svg0).convert_uri_to_datauri()
-    _handlers.config.value = f'<img src="{d.uri}" alt="non-default config">'
-    svg1 = os.path.join(output_dir, 'all-config.svg')
-    save_config_svg(args, svg1)
-    d = Document(uri=svg1).convert_uri_to_datauri()
-    _handlers.all_config.value = f'<img src="{d.uri}" alt="all config">'
+    with tempfile.NamedTemporaryFile(mode='wt', suffix='.svg') as fp:
+        save_config_svg(args, fp.name, only_non_default=True)
+        d = Document(uri=fp.name).convert_uri_to_datauri()
+        _handlers.config.value = f'<img src="{d.uri}" alt="non-default config">'
+
+    with tempfile.NamedTemporaryFile(mode='wt', suffix='.svg') as fp:
+        save_config_svg(args, fp.name)
+        d = Document(uri=fp.name).convert_uri_to_datauri()
+        _handlers.all_config.value = f'<img src="{d.uri}" alt="all config">'
+
     _handlers.code.value = export_python(args)
     _redraw_fn()