[feat/cogstudio-package] Add Cogstudio for CogVideo as package

README.md

@@ -31,6 +31,10 @@
 [sdfx]: https://github.com/sdfxai/sdfx
 [fooocus-mashb1t]: https://github.com/mashb1t/Fooocus
 [reforge]: https://github.com/Panchovix/stable-diffusion-webui-reForge
+[simplesdxl]: https://github.com/metercai/SimpleSDXL/
+[fluxgym]: https://github.com/cocktailpeanut/fluxgym
+[cogvideo]: https://github.com/THUDM/CogVideo
+[cogstudio]: https://github.com/pinokiofactory/cogstudio
 
 [civitai]: https://civitai.com/
 [huggingface]: https://huggingface.co/
@@ -47,15 +51,16 @@ Multi-Platform Package Manager and Inference UI for Stable Diffusion
 ### 🖱️ One click install and update for Stable Diffusion Web UI Packages
 - Supports:
   - [Stable Diffusion WebUI reForge][reforge], [Stable Diffusion WebUI Forge][forge], [Automatic 1111][auto1111], [Automatic 1111 DirectML][auto1111-directml], [SD Web UI-UX][webui-ux], [SD.Next][sdnext]
-  - [Fooocus][fooocus], [Fooocus MRE][fooocus-mre], [Fooocus ControlNet SDXL][fooocus-controlnet], [Ruined Fooocus][ruined-fooocus], [Fooocus - mashb1t's 1-Up Edition][fooocus-mashb1t], [SimpleSDXL](https://github.com/metercai/SimpleSDXL/)
+  - [Fooocus][fooocus], [Fooocus MRE][fooocus-mre], [Fooocus ControlNet SDXL][fooocus-controlnet], [Ruined Fooocus][ruined-fooocus], [Fooocus - mashb1t's 1-Up Edition][fooocus-mashb1t], [SimpleSDXL][simplesdxl]
   - [ComfyUI][comfy]
   - [StableSwarmUI][stable-swarm]
   - [VoltaML][voltaml]
   - [InvokeAI][invokeai]
   - [SDFX][sdfx]
   - [Kohya's GUI][kohya-ss]
   - [OneTrainer][onetrainer]
-  - [FluxGym](https://github.com/cocktailpeanut/fluxgym)
+  - [FluxGym][fluxgym]
+  - [CogVideo][cogvideo] via [CogStudio][cogstudio]
 - Manage plugins / extensions for supported packages ([Automatic1111][auto1111], [Comfy UI][comfy], [SD Web UI-UX][webui-ux], and [SD.Next][sdnext])
 - Easily install or update Python dependencies for each package
 - Embedded Git and Python dependencies, with no need for either to be globally installed

StabilityMatrix.Core/Helper/Factory/PackageFactory.cs

@@ -92,6 +92,8 @@ public BasePackage GetNewBasePackage(InstalledPackage installedPackage)
             "FluxGym" => new FluxGym(githubApiCache, settingsManager, downloadService, prerequisiteHelper),
             "SimpleSDXL"
                 => new SimpleSDXL(githubApiCache, settingsManager, downloadService, prerequisiteHelper),
+            "Cogstudio"
+                => new Cogstudio(githubApiCache, settingsManager, downloadService, prerequisiteHelper),
             _ => throw new ArgumentOutOfRangeException(nameof(installedPackage))
         };
     }

StabilityMatrix.Core/Models/Packages/Cogstudio.cs

@@ -0,0 +1,163 @@
+using System.Text.RegularExpressions;
+using StabilityMatrix.Core.Attributes;
+using StabilityMatrix.Core.Helper;
+using StabilityMatrix.Core.Helper.Cache;
+using StabilityMatrix.Core.Models.FileInterfaces;
+using StabilityMatrix.Core.Models.Progress;
+using StabilityMatrix.Core.Processes;
+using StabilityMatrix.Core.Python;
+using StabilityMatrix.Core.Services;
+
+namespace StabilityMatrix.Core.Models.Packages;
+
+[Singleton(typeof(BasePackage))]
+public class Cogstudio(
+    IGithubApiCache githubApi,
+    ISettingsManager settingsManager,
+    IDownloadService downloadService,
+    IPrerequisiteHelper prerequisiteHelper
+)
+    : BaseGitPackage(githubApi, settingsManager, downloadService, prerequisiteHelper),
+        ISharedFolderLayoutPackage
+{
+    public override string Name => "Cogstudio";
+    public override string DisplayName { get; set; } = "Cogstudio";
+    public override string RepositoryName => "CogVideo";
+    public override string RepositoryAuthor => "THUDM";
+    public override string Author => "pinokiofactory";
+    public override string Blurb =>
+        "An advanced gradio web ui for generating and editing videos with CogVideo.";
+    public override string LicenseType => "Apache-2.0";
+    public override string LicenseUrl => "https://github.com/THUDM/CogVideo/blob/main/LICENSE";
+    public override string LaunchCommand => "inference/gradio_composite_demo/cogstudio.py";
+    public override Uri PreviewImageUri =>
+        new("https://raw.githubusercontent.com/pinokiofactory/cogstudio/main/img2vid.gif");
+    public override List<LaunchOptionDefinition> LaunchOptions => new() { LaunchOptionDefinition.Extras };
+    public override SharedFolderMethod RecommendedSharedFolderMethod => SharedFolderMethod.None;
+    public override IEnumerable<SharedFolderMethod> AvailableSharedFolderMethods =>
+        new[] { SharedFolderMethod.None };
+    public override Dictionary<SharedFolderType, IReadOnlyList<string>> SharedFolders =>
+        ((ISharedFolderLayoutPackage)this).LegacySharedFolders;
+    public virtual SharedFolderLayout SharedFolderLayout => new();
+    public override Dictionary<SharedOutputType, IReadOnlyList<string>> SharedOutputFolders =>
+        new() { [SharedOutputType.Text2Vid] = new[] { "output" } };
+    public override IEnumerable<TorchIndex> AvailableTorchIndices =>
+        new[] { TorchIndex.Cpu, TorchIndex.Cuda };
+    public override string MainBranch => "main";
+    public override bool ShouldIgnoreReleases => true;
+    public override string OutputFolderName => "output";
+    public override PackageDifficulty InstallerSortOrder => PackageDifficulty.Simple;
+
+    public override async Task InstallPackage(
+        string installLocation,
+        InstalledPackage installedPackage,
+        InstallPackageOptions options,
+        IProgress<ProgressReport>? progress = null,
+        Action<ProcessOutput>? onConsoleOutput = null,
+        CancellationToken cancellationToken = default
+    )
+    {
+        const string cogstudioUrl =
+            "https://raw.githubusercontent.com/pinokiofactory/cogstudio/refs/heads/main/cogstudio.py";
+
+        progress?.Report(new ProgressReport(-1f, "Setting up venv", isIndeterminate: true));
+        await using var venvRunner = await SetupVenvPure(installLocation).ConfigureAwait(false);
+
+        progress?.Report(new ProgressReport(-1f, "Setting up Cogstudio files", isIndeterminate: true));
+        var gradioCompositeDemo = new FilePath(installLocation, "inference/gradio_composite_demo");
+        var cogstudioFile = new FilePath(gradioCompositeDemo, "cogstudio.py");
+        gradioCompositeDemo.Directory?.Create();
+        await DownloadService
+            .DownloadToFileAsync(cogstudioUrl, cogstudioFile, cancellationToken: cancellationToken)
+            .ConfigureAwait(false);
+
+        progress?.Report(
+            new ProgressReport(
+                -1f,
+                "Patching cogstudio.py to allow writing to the output folder",
+                isIndeterminate: true
+            )
+        );
+        var outputDir = new FilePath(installLocation, "output");
+        if (Compat.IsWindows)
+        {
+            outputDir = outputDir.ToString().Replace("\\", "\\\\");
+        }
+        var cogstudioContent = await cogstudioFile.ReadAllTextAsync(cancellationToken).ConfigureAwait(false);
+        cogstudioContent = cogstudioContent.Replace(
+            "demo.launch()",
+            $"demo.launch(allowed_paths=['{outputDir}'])"
+        );
+        await cogstudioFile.WriteAllTextAsync(cogstudioContent, cancellationToken).ConfigureAwait(false);
+
+        progress?.Report(new ProgressReport(-1f, "Installing requirements", isIndeterminate: true));
+        var requirements = new FilePath(installLocation, "requirements.txt");
+        var pipArgs = new PipInstallArgs()
+            .WithTorch("==2.3.1")
+            .WithTorchVision("==0.18.1")
+            .WithTorchAudio("==2.3.1")
+            .WithTorchExtraIndex("cu121")
+            .WithParsedFromRequirementsTxt(
+                await requirements.ReadAllTextAsync(cancellationToken).ConfigureAwait(false),
+                excludePattern: Compat.IsWindows
+                    ? "torch.*|moviepy.*|SwissArmyTransformer.*"
+                    : "torch.*|moviepy.*"
+            );
+
+        if (installedPackage.PipOverrides != null)
+        {
+            pipArgs = pipArgs.WithUserOverrides(installedPackage.PipOverrides);
+        }
+
+        // SwissArmyTransformer is not available on Windows and DeepSpeed needs prebuilt wheels
+        if (Compat.IsWindows)
+        {
+            await venvRunner
+                .PipInstall(
+                    " https://github.com/daswer123/deepspeed-windows/releases/download/11.2/deepspeed-0.11.2+cuda121-cp310-cp310-win_amd64.whl",
+                    onConsoleOutput
+                )
+                .ConfigureAwait(false);
+            await venvRunner
+                .PipInstall("spandrel opencv-python scikit-video", onConsoleOutput)
+                .ConfigureAwait(false);
+        }
+
+        await venvRunner.PipInstall(pipArgs, onConsoleOutput).ConfigureAwait(false);
+        await venvRunner.PipInstall("moviepy==2.0.0.dev2", onConsoleOutput).ConfigureAwait(false);
+    }
+
+    public override async Task RunPackage(
+        string installLocation,
+        InstalledPackage installedPackage,
+        RunPackageOptions options,
+        Action<ProcessOutput>? onConsoleOutput = null,
+        CancellationToken cancellationToken = default
+    )
+    {
+        await SetupVenv(installLocation).ConfigureAwait(false);
+
+        void HandleConsoleOutput(ProcessOutput s)
+        {
+            onConsoleOutput?.Invoke(s);
+
+            if (s.Text.Contains("Running on local URL", StringComparison.OrdinalIgnoreCase))
+            {
+                var regex = new Regex(@"(https?:\/\/)([^:\s]+):(\d+)");
+                var match = regex.Match(s.Text);
+
+                if (match.Success)
+                {
+                    WebUrl = match.Value;
+                }
+                OnStartupComplete(WebUrl);
+            }
+        }
+
+        VenvRunner.RunDetached(
+            [Path.Combine(installLocation, options.Command ?? LaunchCommand), ..options.Arguments],
+            HandleConsoleOutput,
+            OnExit
+        );
+    }
+}

StabilityMatrix.Core/Models/SharedOutputType.cs

@@ -4,6 +4,7 @@ public enum SharedOutputType
 {
     All,
     Text2Img,
+    Text2Vid,
     Img2Img,
     Extras,
     Text2ImgGrids,