|
| 1 | +using System.Text.RegularExpressions; |
| 2 | +using StabilityMatrix.Core.Attributes; |
| 3 | +using StabilityMatrix.Core.Helper; |
| 4 | +using StabilityMatrix.Core.Helper.Cache; |
| 5 | +using StabilityMatrix.Core.Models.FileInterfaces; |
| 6 | +using StabilityMatrix.Core.Models.Progress; |
| 7 | +using StabilityMatrix.Core.Processes; |
| 8 | +using StabilityMatrix.Core.Python; |
| 9 | +using StabilityMatrix.Core.Services; |
| 10 | + |
| 11 | +namespace StabilityMatrix.Core.Models.Packages; |
| 12 | + |
| 13 | +[Singleton(typeof(BasePackage))] |
| 14 | +public class Cogstudio( |
| 15 | + IGithubApiCache githubApi, |
| 16 | + ISettingsManager settingsManager, |
| 17 | + IDownloadService downloadService, |
| 18 | + IPrerequisiteHelper prerequisiteHelper |
| 19 | +) |
| 20 | + : BaseGitPackage(githubApi, settingsManager, downloadService, prerequisiteHelper), |
| 21 | + ISharedFolderLayoutPackage |
| 22 | +{ |
| 23 | + public override string Name => "Cogstudio"; |
| 24 | + public override string DisplayName { get; set; } = "Cogstudio"; |
| 25 | + public override string RepositoryName => "CogVideo"; |
| 26 | + public override string RepositoryAuthor => "THUDM"; |
| 27 | + public override string Author => "pinokiofactory"; |
| 28 | + public override string Blurb => |
| 29 | + "An advanced gradio web ui for generating and editing videos with CogVideo."; |
| 30 | + public override string LicenseType => "Apache-2.0"; |
| 31 | + public override string LicenseUrl => "https://github.com/THUDM/CogVideo/blob/main/LICENSE"; |
| 32 | + public override string LaunchCommand => "inference/gradio_composite_demo/cogstudio.py"; |
| 33 | + public override Uri PreviewImageUri => |
| 34 | + new("https://raw.githubusercontent.com/pinokiofactory/cogstudio/main/img2vid.gif"); |
| 35 | + public override List<LaunchOptionDefinition> LaunchOptions => new() { LaunchOptionDefinition.Extras }; |
| 36 | + public override SharedFolderMethod RecommendedSharedFolderMethod => SharedFolderMethod.None; |
| 37 | + public override IEnumerable<SharedFolderMethod> AvailableSharedFolderMethods => |
| 38 | + new[] { SharedFolderMethod.None }; |
| 39 | + public override Dictionary<SharedFolderType, IReadOnlyList<string>> SharedFolders => |
| 40 | + ((ISharedFolderLayoutPackage)this).LegacySharedFolders; |
| 41 | + public virtual SharedFolderLayout SharedFolderLayout => new(); |
| 42 | + public override Dictionary<SharedOutputType, IReadOnlyList<string>> SharedOutputFolders => |
| 43 | + new() { [SharedOutputType.Text2Vid] = new[] { "output" } }; |
| 44 | + public override IEnumerable<TorchIndex> AvailableTorchIndices => |
| 45 | + new[] { TorchIndex.Cpu, TorchIndex.Cuda }; |
| 46 | + public override string MainBranch => "main"; |
| 47 | + public override bool ShouldIgnoreReleases => true; |
| 48 | + public override string OutputFolderName => "output"; |
| 49 | + public override PackageDifficulty InstallerSortOrder => PackageDifficulty.Simple; |
| 50 | + |
| 51 | + public override async Task InstallPackage( |
| 52 | + string installLocation, |
| 53 | + InstalledPackage installedPackage, |
| 54 | + InstallPackageOptions options, |
| 55 | + IProgress<ProgressReport>? progress = null, |
| 56 | + Action<ProcessOutput>? onConsoleOutput = null, |
| 57 | + CancellationToken cancellationToken = default |
| 58 | + ) |
| 59 | + { |
| 60 | + const string cogstudioUrl = |
| 61 | + "https://raw.githubusercontent.com/pinokiofactory/cogstudio/refs/heads/main/cogstudio.py"; |
| 62 | + |
| 63 | + progress?.Report(new ProgressReport(-1f, "Setting up venv", isIndeterminate: true)); |
| 64 | + await using var venvRunner = await SetupVenvPure(installLocation).ConfigureAwait(false); |
| 65 | + |
| 66 | + progress?.Report(new ProgressReport(-1f, "Setting up Cogstudio files", isIndeterminate: true)); |
| 67 | + var gradioCompositeDemo = new FilePath(installLocation, "inference/gradio_composite_demo"); |
| 68 | + var cogstudioFile = new FilePath(gradioCompositeDemo, "cogstudio.py"); |
| 69 | + gradioCompositeDemo.Directory?.Create(); |
| 70 | + await DownloadService |
| 71 | + .DownloadToFileAsync(cogstudioUrl, cogstudioFile, cancellationToken: cancellationToken) |
| 72 | + .ConfigureAwait(false); |
| 73 | + |
| 74 | + progress?.Report( |
| 75 | + new ProgressReport( |
| 76 | + -1f, |
| 77 | + "Patching cogstudio.py to allow writing to the output folder", |
| 78 | + isIndeterminate: true |
| 79 | + ) |
| 80 | + ); |
| 81 | + var outputDir = new FilePath(installLocation, "output"); |
| 82 | + if (Compat.IsWindows) |
| 83 | + { |
| 84 | + outputDir = outputDir.ToString().Replace("\\", "\\\\"); |
| 85 | + } |
| 86 | + var cogstudioContent = await cogstudioFile.ReadAllTextAsync(cancellationToken).ConfigureAwait(false); |
| 87 | + cogstudioContent = cogstudioContent.Replace( |
| 88 | + "demo.launch()", |
| 89 | + $"demo.launch(allowed_paths=['{outputDir}'])" |
| 90 | + ); |
| 91 | + await cogstudioFile.WriteAllTextAsync(cogstudioContent, cancellationToken).ConfigureAwait(false); |
| 92 | + |
| 93 | + progress?.Report(new ProgressReport(-1f, "Installing requirements", isIndeterminate: true)); |
| 94 | + var requirements = new FilePath(installLocation, "requirements.txt"); |
| 95 | + var pipArgs = new PipInstallArgs() |
| 96 | + .WithTorch("==2.3.1") |
| 97 | + .WithTorchVision("==0.18.1") |
| 98 | + .WithTorchAudio("==2.3.1") |
| 99 | + .WithTorchExtraIndex("cu121") |
| 100 | + .WithParsedFromRequirementsTxt( |
| 101 | + await requirements.ReadAllTextAsync(cancellationToken).ConfigureAwait(false), |
| 102 | + excludePattern: Compat.IsWindows |
| 103 | + ? "torch.*|moviepy.*|SwissArmyTransformer.*" |
| 104 | + : "torch.*|moviepy.*" |
| 105 | + ); |
| 106 | + |
| 107 | + if (installedPackage.PipOverrides != null) |
| 108 | + { |
| 109 | + pipArgs = pipArgs.WithUserOverrides(installedPackage.PipOverrides); |
| 110 | + } |
| 111 | + |
| 112 | + // SwissArmyTransformer is not available on Windows and DeepSpeed needs prebuilt wheels |
| 113 | + if (Compat.IsWindows) |
| 114 | + { |
| 115 | + await venvRunner |
| 116 | + .PipInstall( |
| 117 | + " https://github.com/daswer123/deepspeed-windows/releases/download/11.2/deepspeed-0.11.2+cuda121-cp310-cp310-win_amd64.whl", |
| 118 | + onConsoleOutput |
| 119 | + ) |
| 120 | + .ConfigureAwait(false); |
| 121 | + await venvRunner |
| 122 | + .PipInstall("spandrel opencv-python scikit-video", onConsoleOutput) |
| 123 | + .ConfigureAwait(false); |
| 124 | + } |
| 125 | + |
| 126 | + await venvRunner.PipInstall(pipArgs, onConsoleOutput).ConfigureAwait(false); |
| 127 | + await venvRunner.PipInstall("moviepy==2.0.0.dev2", onConsoleOutput).ConfigureAwait(false); |
| 128 | + } |
| 129 | + |
| 130 | + public override async Task RunPackage( |
| 131 | + string installLocation, |
| 132 | + InstalledPackage installedPackage, |
| 133 | + RunPackageOptions options, |
| 134 | + Action<ProcessOutput>? onConsoleOutput = null, |
| 135 | + CancellationToken cancellationToken = default |
| 136 | + ) |
| 137 | + { |
| 138 | + await SetupVenv(installLocation).ConfigureAwait(false); |
| 139 | + |
| 140 | + void HandleConsoleOutput(ProcessOutput s) |
| 141 | + { |
| 142 | + onConsoleOutput?.Invoke(s); |
| 143 | + |
| 144 | + if (s.Text.Contains("Running on local URL", StringComparison.OrdinalIgnoreCase)) |
| 145 | + { |
| 146 | + var regex = new Regex(@"(https?:\/\/)([^:\s]+):(\d+)"); |
| 147 | + var match = regex.Match(s.Text); |
| 148 | + |
| 149 | + if (match.Success) |
| 150 | + { |
| 151 | + WebUrl = match.Value; |
| 152 | + } |
| 153 | + OnStartupComplete(WebUrl); |
| 154 | + } |
| 155 | + } |
| 156 | + |
| 157 | + VenvRunner.RunDetached( |
| 158 | + [Path.Combine(installLocation, options.Command ?? LaunchCommand), ..options.Arguments], |
| 159 | + HandleConsoleOutput, |
| 160 | + OnExit |
| 161 | + ); |
| 162 | + } |
| 163 | +} |
0 commit comments