Merge pull request #53 from axodox/sdxl

Sdxl

Unpaint/ImportHuggingFaceModelDialog.xaml

@@ -17,7 +17,7 @@
   </ContentDialog.Resources>
 
   <StackPanel Spacing="12">
-    <TextBlock TextWrapping="Wrap">You can import Stable Diffusion models from <Hyperlink NavigateUri="https://huggingface.co/models?library=onnx&amp;other=stable-diffusion&amp;other=unpaint">HuggingFace.co</Hyperlink> below. Please note that only ONNX models are supported by Unpaint.</TextBlock>
+    <TextBlock TextWrapping="Wrap">You can import Stable Diffusion models from <Hyperlink NavigateUri="https://huggingface.co/models?library=onnx&amp;other=unpaint">HuggingFace.co</Hyperlink> below. Please note that only ONNX models are supported by Unpaint.</TextBlock>
 
     <Grid ColumnSpacing="12">
       <Grid.ColumnDefinitions>

Unpaint/InferenceOptionsView.xaml

@@ -29,6 +29,7 @@
       <RowDefinition Height="Auto"/>
       <RowDefinition Height="Auto"/>
       <RowDefinition Height="Auto"/>
+      <RowDefinition Height="Auto"/>
     </Grid.RowDefinitions>
 
     <Grid.ColumnDefinitions>
@@ -59,10 +60,15 @@
               Click="{x:Bind ViewModel.ManageModels}"/>
     </Grid>
 
-
     <!-- Resolution -->
-    <TextBlock Grid.Row="2" Text="Resolution" Style="{StaticResource SettingTextBlockStyle}" ToolTipService.ToolTip="Controls the size of the output image. Using higher resolutions will increase memory usage and reduce image generation speed."/>
+    <TextBlock Grid.Row="2" Text="Scheduler" Style="{StaticResource SettingTextBlockStyle}" ToolTipService.ToolTip="Controls how the image is sampled, which affect image quality and aesthetics."/>
     <ComboBox Grid.Row="2" Grid.Column="1" HorizontalAlignment="Stretch"
+          ItemsSource="{x:Bind ViewModel.Schedulers, Mode=OneWay}"
+          SelectedIndex="{x:Bind ViewModel.SelectedSchedulerIndex, Mode=TwoWay}"/>
+
+    <!-- Resolution -->
+    <TextBlock Grid.Row="3" Text="Resolution" Style="{StaticResource SettingTextBlockStyle}" ToolTipService.ToolTip="Controls the size of the output image. Using higher resolutions will increase memory usage and reduce image generation speed."/>
+    <ComboBox Grid.Row="3" Grid.Column="1" HorizontalAlignment="Stretch"
               ItemsSource="{x:Bind ViewModel.Resolutions, Mode=OneWay}"
               SelectedIndex="{x:Bind ViewModel.SelectedResolutionIndex, Mode=TwoWay}">
       <ComboBox.ItemTemplate>
@@ -73,29 +79,29 @@
     </ComboBox>
 
     <!-- Batch size -->
-    <CheckBox Grid.Row="3" Content="Batching" VerticalAlignment="Center" HorizontalAlignment="Right" FlowDirection="RightToLeft"
+    <CheckBox Grid.Row="4" Content="Batching" VerticalAlignment="Center" HorizontalAlignment="Right" FlowDirection="RightToLeft"
               IsChecked="{x:Bind ViewModel.IsBatchGenerationEnabled, Mode=TwoWay}">
       <ToolTipService.ToolTip>
         <TextBlock Text="Generate multiple images in one run: this can increase per image efficiency significantly, however generation will take longer and have increased video memory usage. The efficiency improvement diminishes at higher values." TextWrapping="WrapWholeWords" FlowDirection="LeftToRight"/>
       </ToolTipService.ToolTip>
     </CheckBox>
-    <muxc:NumberBox Grid.Row="3" Grid.Column="1" Minimum="1" Maximum="128" SmallChange="1" SpinButtonPlacementMode="Inline"
+    <muxc:NumberBox Grid.Row="4" Grid.Column="1" Minimum="1" Maximum="128" SmallChange="1" SpinButtonPlacementMode="Inline"
                     IsEnabled="{x:Bind ViewModel.IsBatchGenerationEnabled, Mode=OneWay}"
                     Value="{x:Bind ViewModel.BatchSize, Mode=TwoWay}"/>
 
     <!-- Guidance scale -->
-    <TextBlock Grid.Row="4" Text="Guidance strength" Style="{StaticResource SettingTextBlockStyle}" ToolTipService.ToolTip="Controls the strength of human input when forming the output. Low values can result in more creative results, while higher values will conform to the input more strictly. Selecting values too high or low degrade the quality of the output."/>
-    <Slider Grid.Row="4" Grid.Column="1" Minimum="0" Maximum="20" SmallChange="0.1" LargeChange="1" StepFrequency="0.1"
+    <TextBlock Grid.Row="5" Text="Guidance strength" Style="{StaticResource SettingTextBlockStyle}" ToolTipService.ToolTip="Controls the strength of human input when forming the output. Low values can result in more creative results, while higher values will conform to the input more strictly. Selecting values too high or low degrade the quality of the output."/>
+    <Slider Grid.Row="5" Grid.Column="1" Minimum="0" Maximum="20" SmallChange="0.1" LargeChange="1" StepFrequency="0.1"
             Value="{x:Bind ViewModel.GuidanceStrength, Mode=TwoWay}"/>
 
     <!-- Sampling steps -->
-    <TextBlock Grid.Row="5" Text="Sampling steps" Style="{StaticResource SettingTextBlockStyle}" ToolTipService.ToolTip="Reducing this value speeds up image generation time, however it also degrades image quality."/>
-    <muxc:NumberBox Grid.Row="5" Grid.Column="1" Minimum="5" Maximum="100" SmallChange="1" SpinButtonPlacementMode="Inline"
+    <TextBlock Grid.Row="6" Text="Sampling steps" Style="{StaticResource SettingTextBlockStyle}" ToolTipService.ToolTip="Reducing this value speeds up image generation time, however it also degrades image quality."/>
+    <muxc:NumberBox Grid.Row="6" Grid.Column="1" Minimum="5" Maximum="100" SmallChange="1" SpinButtonPlacementMode="Inline"
                     Value="{x:Bind ViewModel.SamplingSteps, Mode=TwoWay}"/>
 
     <!-- Random seed -->
-    <TextBlock Grid.Row="6" Text="Random seed" Style="{StaticResource SettingTextBlockStyle}" ToolTipService.ToolTip="Affects random numbers used in the image generation process. Using the same seed number with the same settings will result in the same image being generated. Useful for testing how different settings affect the output image."/>
-    <Grid Grid.Row="6" Grid.Column="1" ColumnSpacing="6">
+    <TextBlock Grid.Row="7" Text="Random seed" Style="{StaticResource SettingTextBlockStyle}" ToolTipService.ToolTip="Affects random numbers used in the image generation process. Using the same seed number with the same settings will result in the same image being generated. Useful for testing how different settings affect the output image."/>
+    <Grid Grid.Row="7" Grid.Column="1" ColumnSpacing="6">
       <Grid.ColumnDefinitions>
         <ColumnDefinition Width="*"/>
         <ColumnDefinition Width="Auto"/>

Unpaint/InferenceOptionsViewModel.cpp

@@ -3,6 +3,7 @@
 #include "InferenceOptionsViewModel.g.cpp"
 
 using namespace Axodox::Infrastructure;
+using namespace Axodox::MachineLearning;
 using namespace winrt::Windows::Graphics;
 using namespace winrt::Windows::UI::Xaml::Data;
 
@@ -14,6 +15,7 @@ namespace winrt::Unpaint::implementation
     _modelRepository(dependencies.resolve<ModelRepository>()),
     _deviceInformation(dependencies.resolve<DeviceInformation>()),
     _models(single_threaded_observable_vector<ModelViewModel>()),
+    _schedulers(single_threaded_observable_vector<hstring>()),
     _resolutions(single_threaded_observable_vector<SizeInt32>()),
     _selectedResolutionIndex(1),
     _selectedModelIndex(0),
@@ -22,6 +24,10 @@ namespace winrt::Unpaint::implementation
     //Initialize models
     OnModelChanged();
 
+    //Initialize schedulers
+    _schedulers.Append(L"Euler A");
+    _schedulers.Append(L"DPM++ 2M Karras");
+
     //Initialize resolutions
     _resolutions.Append(SizeInt32{ 1024, 1024 });
     _resolutions.Append(SizeInt32{ 768, 768 });
@@ -48,6 +54,24 @@ namespace winrt::Unpaint::implementation
     _propertyChanged(*this, PropertyChangedEventArgs(L"SelectedModelIndex"));
   }
 
+  winrt::Windows::Foundation::Collections::IObservableVector<winrt::hstring> InferenceOptionsViewModel::Schedulers()
+  {
+    return _schedulers;
+  }
+
+  int32_t InferenceOptionsViewModel::SelectedSchedulerIndex()
+  {
+    return int32_t(*_unpaintState->Scheduler);
+  }
+
+  void InferenceOptionsViewModel::SelectedSchedulerIndex(int32_t value)
+  {
+    if (value == SelectedSchedulerIndex()) return;
+
+    _unpaintState->Scheduler = StableDiffusionSchedulerKind(value);
+    _propertyChanged(*this, PropertyChangedEventArgs(L"SelectedSchedulerIndex"));
+  }
+
   winrt::Windows::Foundation::Collections::IObservableVector<winrt::Windows::Graphics::SizeInt32> InferenceOptionsViewModel::Resolutions()
   {
     return _resolutions;

Unpaint/InferenceOptionsViewModel.h

@@ -14,6 +14,11 @@ namespace winrt::Unpaint::implementation
     int32_t SelectedModelIndex();
     void SelectedModelIndex(int32_t value);
 
+    winrt::Windows::Foundation::Collections::IObservableVector<winrt::hstring> Schedulers();
+
+    int32_t SelectedSchedulerIndex();
+    void SelectedSchedulerIndex(int32_t value);
+
     winrt::Windows::Foundation::Collections::IObservableVector<winrt::Windows::Graphics::SizeInt32> Resolutions();
 
     int32_t SelectedResolutionIndex();
@@ -51,6 +56,8 @@ namespace winrt::Unpaint::implementation
     Windows::Foundation::Collections::IObservableVector<ModelViewModel> _models;
     int32_t _selectedModelIndex;
 
+    Windows::Foundation::Collections::IObservableVector<hstring> _schedulers;
+
     Windows::Foundation::Collections::IObservableVector<Windows::Graphics::SizeInt32> _resolutions;
     int32_t _selectedResolutionIndex;
 

Unpaint/InferenceView.xaml

@@ -333,7 +333,7 @@
         </Grid>
 
         <!-- Mode selector -->
-        <ComboBox Grid.Column="1" Width="120" SelectedIndex="{x:Bind ViewModel.SelectedModeIndex, Mode=TwoWay}">
+        <ComboBox Grid.Column="1" Width="120" SelectedIndex="{x:Bind ViewModel.SelectedModeIndex, Mode=TwoWay}" IsEnabled="{x:Bind ViewModel.IsModeSelectable, Mode=OneWay}">
           <ToolTipService.ToolTip>
             <StackPanel>
               <TextBlock Text="Mode Selector" FontWeight="Bold"/>

Unpaint/InferenceViewModel.cpp

@@ -53,7 +53,8 @@ namespace winrt::Unpaint::implementation
     _featureMask(nullptr),
     _inputResolution({ 0, 0 }),
     _isAutoGenerationEnabled(false),
-    _safetyStrikes(0)
+    _safetyStrikes(0),
+    _modelChangedSubscription(_unpaintState->ModelId.ValueChanged(event_handler{ this, &InferenceViewModel::OnModelChanged }))
   { }
 
   ProjectViewModel InferenceViewModel::Project()
@@ -74,6 +75,12 @@ namespace winrt::Unpaint::implementation
     _propertyChanged(*this, PropertyChangedEventArgs(L"SelectedModeIndex"));
   }
 
+  bool InferenceViewModel::IsModeSelectable()
+  {
+    auto model = _modelRepository->GetModel(*_unpaintState->ModelId);
+    return model && !model->IsXL;
+  }
+
   bool InferenceViewModel::IsSettingsLocked()
   {
     return _unpaintState->IsSettingsLocked;
@@ -250,6 +257,7 @@ namespace winrt::Unpaint::implementation
       .Mode = _unpaintState->InferenceMode,
       .PositivePrompt = _unpaintState->PositivePrompt->empty() ? StableDiffusionInferenceTask::PositivePromptPlaceholder : *_unpaintState->PositivePrompt,
       .NegativePrompt = _unpaintState->NegativePrompt->empty() ? StableDiffusionInferenceTask::NegativePromptPlaceholder : *_unpaintState->NegativePrompt,
+      .Scheduler = _unpaintState->Scheduler,
       .Resolution = { uint32_t(_unpaintState->Resolution->Width), uint32_t(_unpaintState->Resolution->Height) },
       .GuidanceStrength = _unpaintState->GuidanceStrength,
       .DenoisingStrength = _unpaintState->DenoisingStrength,
@@ -258,7 +266,7 @@ namespace winrt::Unpaint::implementation
       .BatchSize = _unpaintState->IsBatchGenerationEnabled ? *_unpaintState->BatchSize : 1,
       .IsSafeModeEnabled = _unpaintState->IsSafeModeEnabled,
       .IsSafetyCheckerEnabled = _unpaintState->IsSafetyCheckerEnabled,
-      .ModelId = _unpaintState->ModelId      
+      .ModelId = _unpaintState->ModelId
     };
 
     filesystem::path inputPath;
@@ -461,4 +469,14 @@ namespace winrt::Unpaint::implementation
   {
     _propertyChanged.remove(token);
   }
+
+  void InferenceViewModel::OnModelChanged(OptionPropertyBase* /*option*/)
+  {
+    if (!IsModeSelectable())
+    {
+      SelectedModeIndex(0);
+    }
+
+    _propertyChanged(*this, PropertyChangedEventArgs(L"IsModeSelectable"));
+  }
 }

Unpaint/InferenceViewModel.h

@@ -19,6 +19,8 @@ namespace winrt::Unpaint::implementation
     int32_t SelectedModeIndex();
     void SelectedModeIndex(int32_t value);
 
+    bool IsModeSelectable();
+
     bool IsSettingsLocked();
     void IsSettingsLocked(bool value);
 
@@ -87,6 +89,10 @@ namespace winrt::Unpaint::implementation
     bool _isAutoGenerationEnabled;
     bool _hasSafetyCheckFailed;
     uint32_t _safetyStrikes;
+
+    Axodox::Infrastructure::event_subscription _modelChangedSubscription;
+
+    void OnModelChanged(OptionPropertyBase* option);
   };
 }
 

Unpaint/ModelRepository.cpp

@@ -122,11 +122,14 @@ namespace winrt::Unpaint
       auto metadata = try_parse_json<ModelMetadata>(*text);
       if (!metadata) continue;
 
+      auto isXL = filesystem::exists(file.path().parent_path() / "text_encoder_2", ec);
+
       models.emplace(ModelInfo{ 
         *metadata->Id, 
         metadata->Name->empty() ? *metadata->Id : *metadata->Name,
         *metadata->Website,
-        *metadata->AccessToken 
+        *metadata->AccessToken,
+        isXL
       });
     }
 
@@ -192,7 +195,8 @@ namespace winrt::Unpaint
     return ModelViewModel{
       .Id = to_hstring(Id),
       .Name = to_hstring(Name),
-      .Uri = to_hstring(Website)
+      .Uri = to_hstring(Website),
+      .IsXL = IsXL
     };
   }
 }

Unpaint/ModelRepository.h

@@ -23,6 +23,7 @@ namespace winrt::Unpaint
     std::string Name;
     std::string Website;
     std::string AccessToken;
+    bool IsXL;
 
     auto operator<=>(const ModelInfo&) const = default;
     bool operator<(const ModelInfo&) const = default;

Unpaint/SettingsViewModel.cpp

@@ -30,12 +30,9 @@ namespace winrt::Unpaint::implementation
 
   bool SettingsViewModel::AreUnsafeOptionsEnabled()
   {
-#ifdef NDEBUG
-    //Sorry mates I do not trust you this much...
-    return false;
-#else
-    return true;
-#endif
+#pragma warning(suppress: 4996)
+    auto devMode = getenv("UNPAINT_DEV");
+    return devMode && strcmp(devMode, "1") == 0;
   }
 
   bool SettingsViewModel::IsSafeModeEnabled()

Unpaint/StableDiffusionModelExecutor.cpp

@@ -277,13 +277,13 @@ namespace winrt::Unpaint
     for (auto weight : encodedPositivePrompt[0].Weights) textEmbbedding.Weights.push_back(weight);
 
     ScheduledTensor tensor{ task.SamplingSteps };
-    trivial_map<pair<void*, void*>, shared_ptr<Tensor>> embeddingBuffer;
+    trivial_map<pair<void*, void*>, shared_ptr<EncodedText>> embeddingBuffer;
     for (auto i = 0u; i < task.SamplingSteps; i++)
     {
       auto& concatenatedTensor = embeddingBuffer[{ encodedNegativePrompt[i].Tensor.get(), encodedPositivePrompt[i].Tensor.get() }];
       if (!concatenatedTensor)
       {
-        concatenatedTensor = make_shared<Tensor>(encodedNegativePrompt[i].Tensor->Concat(*encodedPositivePrompt[i].Tensor));
+        concatenatedTensor = make_shared<EncodedText>(encodedNegativePrompt[i].Tensor->Concat(*encodedPositivePrompt[i].Tensor));
       }
 
       tensor[i] = concatenatedTensor;
@@ -330,7 +330,8 @@ namespace winrt::Unpaint
       .TextEmbeddings = inputs.TextEmbeddings,
       .LatentInput = inputs.InputImage,
       .MaskInput = inputs.InputMask,
-      .DenoisingStrength = task.Mode == InferenceMode::Modify ? task.DenoisingStrength : 1.f
+      .DenoisingStrength = task.Mode == InferenceMode::Modify ? task.DenoisingStrength : 1.f,
+      .Scheduler = task.Scheduler
     };
 
     async.update_state("Running denoiser...");
@@ -366,8 +367,11 @@ namespace winrt::Unpaint
 
   void StableDiffusionModelExecutor::RunSafetyCheck(std::vector<Axodox::Graphics::TextureData>& images, Axodox::Threading::async_operation_source& async)
   {
+    static const char* safetyCheckedId = "safety_checker\\model.onnx";
+    if (!_modelFiles.contains(safetyCheckedId)) return;
+
     async.update_state(NAN, "Loading safety checker...");
-    SafetyChecker safetyChecker{ *_onnxEnvironment, GetModelFile("safety_checker\\model.onnx") };
+    SafetyChecker safetyChecker{ *_onnxEnvironment, GetModelFile(safetyCheckedId) };
 
     async.update_state(NAN, "Checking safety...");
     for (auto& image : images)

Unpaint/StableDiffusionModelExecutor.h

@@ -21,6 +21,7 @@ namespace winrt::Unpaint
     InferenceMode Mode;
 
     std::string PositivePrompt, NegativePrompt;
+    Axodox::MachineLearning::StableDiffusionSchedulerKind Scheduler;
     DirectX::XMUINT2 Resolution;
 
     float GuidanceStrength;