buggy with multiple debug statements

Author: Lincoln Stein

invokeai/app/invocations/compel.py

@@ -1,5 +1,5 @@
 from typing import Iterator, List, Optional, Tuple, Union, cast
-
+import threading # for debugging only
 import torch
 from compel import Compel, ReturnedEmbeddingsType
 from compel.prompt_parser import Blend, Conjunction, CrossAttentionControlSubstitute, FlattenedPrompt, Fragment
@@ -91,6 +91,9 @@ def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
                 ti_manager,
             ),
         ):
+
+            print(f'DEBUG: compel: tid={threading.current_thread().ident}, gpu={TorchDevice.choose_torch_device()}, text_encoder={text_encoder.device}')
+
             assert isinstance(text_encoder, CLIPTextModel)
             assert isinstance(tokenizer, CLIPTokenizer)
             compel = Compel(
@@ -99,6 +102,7 @@ def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
                 textual_inversion_manager=ti_manager,
                 dtype_for_device_getter=TorchDevice.choose_torch_dtype,
                 truncate_long_prompts=False,
+                device=TorchDevice.choose_torch_device(),
             )
 
             conjunction = Compel.parse_prompt_string(self.prompt)
@@ -113,6 +117,8 @@ def _lora_loader() -> Iterator[Tuple[LoRAModelRaw, float]]:
         conditioning_data = ConditioningFieldData(conditionings=[BasicConditioningInfo(embeds=c)])
 
         conditioning_name = context.conditioning.save(conditioning_data)
+        print(f'DEBUG: conditioning_name={conditioning_name}')
+
         return ConditioningOutput(
             conditioning=ConditioningField(
                 conditioning_name=conditioning_name,

invokeai/app/invocations/latent.py

@@ -1,6 +1,8 @@
 # Copyright (c) 2023 Kyle Schouviller (https://github.com/kyle0654)
 import inspect
 import math
+import threading  # for debugging
+
 from contextlib import ExitStack
 from functools import singledispatchmethod
 from typing import Any, Dict, Iterator, List, Literal, Optional, Tuple, Union
@@ -1078,6 +1080,7 @@ def invoke(self, context: InvocationContext) -> ImageOutput:
             with torch.inference_mode():
                 # copied from diffusers pipeline
                 latents = latents / vae.config.scaling_factor
+                print(f'DEBUG: tid={threading.current_thread().ident}, gpu={TorchDevice.choose_torch_device()}, latent_device={latents.device}')
                 image = vae.decode(latents, return_dict=False)[0]
                 image = (image / 2 + 0.5).clamp(0, 1)  # denormalize
                 # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16

invokeai/app/services/invocation_stats/invocation_stats_default.py

@@ -74,9 +74,9 @@ def collect_stats(self, invocation: BaseInvocation, graph_execution_state_id: st
             )
             self._stats[graph_execution_state_id].add_node_execution_stats(node_stats)
 
-    def reset_stats(self):
-        self._stats = {}
-        self._cache_stats = {}
+    def reset_stats(self, graph_execution_state_id: str):
+        self._stats.pop(graph_execution_state_id)
+        self._cache_stats.pop(graph_execution_state_id)
 
     def get_stats(self, graph_execution_state_id: str) -> InvocationStatsSummary:
         graph_stats_summary = self._get_graph_summary(graph_execution_state_id)

invokeai/app/services/session_processor/session_processor_default.py

@@ -214,7 +214,7 @@ def _on_after_run_session(self, queue_item: SessionQueueItem) -> None:
             # we don't care about that - suppress the error.
             with suppress(GESStatsNotFoundError):
                 self._services.performance_statistics.log_stats(queue_item.session.id)
-                self._services.performance_statistics.reset_stats()
+                self._services.performance_statistics.reset_stats(queue_item.session.id)
 
             for callback in self._on_after_run_session_callbacks:
                 callback(queue_item=queue_item)

invokeai/app/services/shared/invocation_context.py

@@ -325,7 +325,6 @@ def load(self, name: str) -> ConditioningFieldData:
         Returns:
             The loaded conditioning data.
         """
-
         return self._services.conditioning.load(name)