Merge pull request #1962 from Capsize-Games/develop

improve zimage and add pixtral support

Author: w4ffl35

setup.py

@@ -187,7 +187,7 @@
 
 setup(
     name="airunner",
-    version="5.4.2",
+    version="5.5.0",
     author="Capsize LLC",
     description="Run local opensource AI models (Stable Diffusion, LLMs, TTS, STT, chatbots) in a lightweight Python GUI",
     long_description=open("README.md", "r", encoding="utf-8").read(),

src/airunner/components/art/gui/widgets/lora/lora_container_widget.py

@@ -23,7 +23,7 @@ class LoraContainerWidget(BaseWidget):
 
     def __init__(self, *args, **kwargs):
         self.signal_handlers = {
-            SignalCode.APPLICATION_SETTINGS_CHANGED_SIGNAL: self.on_application_settings_changed_signal,
+            # SignalCode.APPLICATION_SETTINGS_CHANGED_SIGNAL: self.on_application_settings_changed_signal,
             SignalCode.LORA_UPDATED_SIGNAL: self.on_lora_updated_signal,
             SignalCode.MODEL_STATUS_CHANGED_SIGNAL: self.on_model_status_changed_signal,
             SignalCode.LORA_STATUS_CHANGED: self.on_lora_modified,
@@ -115,7 +115,7 @@ def scan_for_lora(self):
         self._load_lora(force_reload=force_reload)
 
     @Slot()
-    def apply_lora(self):
+    def on_apply_lora_button_clicked(self):
         self._apply_button_enabled = False
         self.api.art.lora.update()
 
@@ -197,6 +197,7 @@ def _load_lora(self, force_reload=False):
                     if self.search_filter.lower() in lora.name.lower()
                 ]
                 for lora in filtered_loras:
+                    print("adding lora widget for", lora.name)
                     self._add_lora(lora)
                 self.add_spacer()
 

src/airunner/components/art/gui/widgets/lora/templates/lora_container.ui

@@ -280,8 +280,8 @@
    <slot>toggle_all(bool)</slot>
    <hints>
     <hint type="sourcelabel">
-     <x>562</x>
-     <y>91</y>
+     <x>561</x>
+     <y>95</y>
     </hint>
     <hint type="destinationlabel">
      <x>52</x>
@@ -312,31 +312,15 @@
    <slot>scan_for_lora()</slot>
    <hints>
     <hint type="sourcelabel">
-     <x>399</x>
-     <y>830</y>
+     <x>409</x>
+     <y>820</y>
     </hint>
     <hint type="destinationlabel">
      <x>486</x>
      <y>-1</y>
     </hint>
    </hints>
   </connection>
-  <connection>
-   <sender>apply_lora_button</sender>
-   <signal>clicked()</signal>
-   <receiver>lora_container</receiver>
-   <slot>apply_lora()</slot>
-   <hints>
-    <hint type="sourcelabel">
-     <x>532</x>
-     <y>17</y>
-    </hint>
-    <hint type="destinationlabel">
-     <x>474</x>
-     <y>-7</y>
-    </hint>
-   </hints>
-  </connection>
  </connections>
  <slots>
   <slot>toggle_all(bool)</slot>

src/airunner/components/art/gui/widgets/lora/templates/lora_container_ui.py

@@ -148,7 +148,6 @@ def setupUi(self, lora_container):
         self.toggleAllLora.toggled.connect(lora_container.toggle_all)
         self.lineEdit.textEdited.connect(lora_container.search_text_changed)
         self.pushButton.clicked.connect(lora_container.scan_for_lora)
-        self.apply_lora_button.clicked.connect(lora_container.apply_lora)
 
         QMetaObject.connectSlotsByName(lora_container)
     # setupUi

src/airunner/components/art/managers/stablediffusion/mixins/sd_model_loading_mixin.py

@@ -251,6 +251,7 @@ def _load_lora_weights(self, lora: Lora):
                 lora_base_path,
                 weight_name=filename,
                 adapter_name=adapter_name,
+                scale=getattr(lora, "scale", 1.0),
             )
             self._loaded_lora[lora.path] = lora
         except AttributeError:

src/airunner/components/art/managers/stablediffusion/model_loader.py

@@ -127,11 +127,13 @@ def load_lora_weights(
 ) -> bool:
     filename = os.path.basename(lora.path)
     adapter_name = os.path.splitext(filename)[0].replace(".", "_")
+    # Scale is stored as 0-100 integer, convert to 0.0-1.0 float
+    scale = lora.scale / 100.0 if hasattr(lora, 'scale') else 1.0
     try:
         pipe.load_lora_weights(
-            lora_base_path, weight_name=filename, adapter_name=adapter_name
+            lora_base_path, weight_name=filename, adapter_name=adapter_name, scale=scale
         )
-        logger.info(f"Loaded LORA weights: {filename}")
+        logger.info(f"Loaded LORA weights: {filename} (scale={scale:.2f})")
         return True
     except Exception as e:
         logger.warning(f"Failed to load LORA {filename}: {e}")

src/airunner/components/art/managers/zimage/mixins/zimage_generation_mixin.py

@@ -6,6 +6,13 @@
 from airunner.components.application.exceptions import InterruptedException
 from airunner.utils.memory import clear_memory
 from airunner.utils.settings.get_qsettings import get_qsettings
+from airunner.components.art.schedulers.flow_match_scheduler_factory import (
+    is_flow_match_scheduler,
+    create_flow_match_scheduler,
+)
+from airunner.enums import Scheduler, ModelType, ModelStatus
+
+from accelerate.hooks import remove_hook_from_module
 
 
 def _aggressive_memory_cleanup():
@@ -57,7 +64,6 @@ def _prepare_data(self, active_rect=None) -> Dict:
         self._strip_zimage_incompatible_params(data)
         self._enforce_zimage_guidance(data)
         data["max_sequence_length"] = 512
-        self._log_zimage_generation_params(data)
         return data
 
     def _strip_zimage_incompatible_params(self, data: Dict) -> None:
@@ -81,34 +87,40 @@ def _enforce_zimage_guidance(self, data: Dict) -> None:
         """
         pass
 
-    def _log_zimage_generation_params(self, data: Dict) -> None:
-        """Log core generation parameters for debugging."""
-        debug_fields = {
-            "prompt": data.get("prompt", "MISSING!")[:50] + "...",
-            "guidance_scale": data.get("guidance_scale", "MISSING!"),
-            "steps": data.get("num_inference_steps", "MISSING!"),
-            "size": f"{data.get('width')}x{data.get('height')}",
-            "max_sequence_length": data.get("max_sequence_length", "MISSING!"),
-        }
-        self.logger.info(
-            "[Z-IMAGE DEBUG] Keys: %s | Values: %s",
-            list(data.keys()),
-            debug_fields,
-        )
-
     def _unload_loras(self):
         """Unload Z-Image LoRA weights if any are loaded.
         
-        Z-Image supports LoRA weights through ZImageLoraLoaderMixin.
+        Z-Image uses additive LoRA that can be removed without model reload.
         """
-        if hasattr(self._pipe, 'unload_lora_weights'):
+        self.logger.debug("Unloading Z-Image LoRA weights")
+        if self._pipe is not None and hasattr(self._pipe, 'unload_lora_weights'):
             try:
                 self._pipe.unload_lora_weights()
-                self.logger.debug("Unloaded Z-Image LoRA weights")
+                self.logger.info("✓ Unloaded all Z-Image LoRA weights")
             except Exception as e:
-                self.logger.debug(f"No LoRA weights to unload: {e}")
+                self.logger.warning(f"Error unloading LoRA weights: {e}")
         self._loaded_lora = {}
         self._disabled_lora = []
+    
+    def _disable_lora(self, adapter_name: str):
+        """Disable a specific LoRA adapter without removing it.
+        
+        Args:
+            adapter_name: Name of the adapter to disable
+        """
+        if self._pipe is not None and hasattr(self._pipe, 'set_lora_enabled'):
+            self._pipe.set_lora_enabled(adapter_name, False)
+            self.logger.debug(f"Disabled LoRA: {adapter_name}")
+    
+    def _enable_lora(self, adapter_name: str):
+        """Enable a specific LoRA adapter.
+        
+        Args:
+            adapter_name: Name of the adapter to enable
+        """
+        if self._pipe is not None and hasattr(self._pipe, 'set_lora_enabled'):
+            self._pipe.set_lora_enabled(adapter_name, True)
+            self.logger.debug(f"Enabled LoRA: {adapter_name}")
 
     def _load_scheduler(self, scheduler_name=None):
         """Load a flow-match scheduler for Z-Image.
@@ -118,11 +130,7 @@ def _load_scheduler(self, scheduler_name=None):
         Args:
             scheduler_name: Display name of the scheduler to load.
         """
-        from airunner.components.art.schedulers.flow_match_scheduler_factory import (
-            is_flow_match_scheduler,
-            create_flow_match_scheduler,
-        )
-        from airunner.enums import Scheduler, ModelType, ModelStatus
+        # imports moved to module level for performance and clarity
 
         # Get scheduler name
         requested_name = (
@@ -232,11 +240,8 @@ def _unload_pipe(self):
             return
 
         # Import accelerate hooks removal if available
-        try:
-            from accelerate.hooks import remove_hook_from_module
-            has_accelerate_hooks = True
-        except ImportError:
-            has_accelerate_hooks = False
+        has_accelerate_hooks = remove_hook_from_module is not None
+        if not has_accelerate_hooks:
             self.logger.debug("accelerate.hooks not available, using manual cleanup")
 
         # List of all Z-Image components to clean up (ordered by size)
@@ -259,78 +264,12 @@ def _unload_pipe(self):
                         self.logger.debug(f"Error removing hook: {e}")
                 self._pipe._all_hooks.clear()
 
-            # Process each component
+            # Process each component via helper
             for component_name in component_names:
                 component = getattr(self._pipe, component_name, None)
                 if component is None:
                     continue
-                
-                self.logger.debug(f"Cleaning up {component_name}...")
-                
-                # 1. Remove accelerate hooks using official API
-                if has_accelerate_hooks:
-                    try:
-                        remove_hook_from_module(component, recurse=True)
-                        self.logger.debug(f"Removed hooks from {component_name}")
-                    except Exception as e:
-                        self.logger.debug(f"Hook removal for {component_name}: {e}")
-                
-                # 2. Manual hook cleanup for any remaining hooks
-                if hasattr(component, "_hf_hook"):
-                    try:
-                        hook = component._hf_hook
-                        # Clear any offloaded weights first
-                        if hasattr(hook, "weights_map") and hook.weights_map is not None:
-                            hook.weights_map.clear()
-                        if hasattr(hook, "offload"):
-                            try:
-                                hook.offload(component)
-                            except Exception:
-                                pass
-                        delattr(component, "_hf_hook")
-                    except Exception as e:
-                        self.logger.debug(f"Manual hook cleanup for {component_name}: {e}")
-                
-                # 3. For models with device_map (quantized models), clear the device_map state
-                if hasattr(component, "hf_device_map"):
-                    try:
-                        component.hf_device_map = None
-                    except Exception:
-                        pass
-                
-                # 4. CRITICAL: Do NOT move to CPU - this creates new tensors in CPU RAM
-                # Instead, delete tensors in-place on their current device
-                
-                # 5. Clear parameter data to free memory
-                if hasattr(component, "parameters"):
-                    try:
-                        for param in component.parameters():
-                            param.data = torch.empty(0, device=param.device)
-                            if param.grad is not None:
-                                param.grad = None
-                    except Exception as e:
-                        self.logger.debug(f"Error clearing params for {component_name}: {e}")
-                
-                # 6. Detach the component from the pipeline
-                try:
-                    setattr(self._pipe, component_name, None)
-                except Exception:
-                    pass
-                
-                # 7. Delete component reference
-                try:
-                    del component
-                except Exception:
-                    pass
-                
-                # 8. Run gc after each large component to free memory immediately
-                if component_name in ("text_encoder", "transformer"):
-                    _aggressive_memory_cleanup()
-            
-            # Clear execution device reference
-            if hasattr(self._pipe, "_execution_device"):
-                self._pipe._execution_device = None
-                
+                self._cleanup_pipeline_component(component_name, component, has_accelerate_hooks)
         except Exception as e:
             self.logger.warning(f"Error during hook removal: {e}")
 
@@ -347,20 +286,46 @@ def _unload_pipe(self):
 
         self.logger.info("✓ Z-Image pipeline unloaded and memory freed")
 
-    def _clear_pipeline_caches(self):
-        """Clear internal pipeline caches to free RAM and VRAM.
+    def _cleanup_pipeline_component(self, component_name: str, component: Any, has_accelerate_hooks: bool) -> None:
+        """Cleanup a component attached to the pipeline.
+
+        Extraction from the prior implementation to reduce _unload_pipe size.
+        """
+        self.logger.debug(f"Cleaning up {component_name}...")
+
+        # 1. Remove accelerate hooks using official API
+        if has_accelerate_hooks:
+            try:
+                remove_hook_from_module(component, recurse=True)
+                self.logger.debug(f"Removed hooks from {component_name}")
+            except Exception as e:
+                self.logger.debug(f"Hook removal for {component_name}: {e}")
+
+        # 2. Manual hook cleanup for any remaining hooks
+        if hasattr(component, "_hf_hook"):
+            try:
+                hook = component._hf_hook
+                if hasattr(hook, "weights_map") and hook.weights_map is not None:
+                    hook.weights_map.clear()
+                if hasattr(hook, "offload"):
+                    try:
+                        hook.offload(component)
+                    except Exception:
+                        pass
+                delattr(component, "_hf_hook")
+            except Exception as e:
+                self.logger.debug(f"Manual hook cleanup for {component_name}: {e}")
 
+    def _clear_pipeline_caches(self):
+        """Clear cached tensors and per-component caches on the active pipeline.
+
         This is called after each generation to prevent memory accumulation.
         """
         if self._pipe is None:
             return
 
         self.logger.debug("Clearing pipeline caches to free RAM")
 
-        # Clear any cached tensors on the pipeline
-        if hasattr(self._pipe, "_callback_tensor_inputs"):
-            self._pipe._callback_tensor_inputs = None
-        
         # For text encoder, clear any cached key/values
         text_encoder = getattr(self._pipe, "text_encoder", None)
         if text_encoder is not None and hasattr(text_encoder, "past_key_values"):