Convenience script for basujindal/stable-diffusion

generate.ps1

# Interactive convenience script for basujindal/stable-diffusion.

$Settings = @{
    Prompt = "your prompt here"
    Steps = "60"
    Iters = "3"
    Samples = "1"
    Seed = "12345"
    ProgressiveGeneration = "y"
    UseOptimizedScript = "n"
}

# Load previous settings json if it exists
$SettingsPath = './settings.json'
if (Test-Path $SettingsPath) {
    $Settings = Get-Content $SettingsPath -Raw
    $Settings = ConvertFrom-Json $Settings
    Set-Clipboard $($Settings.Prompt)
    Write-Host "Your last prompt was copied to the clipboard."
}

# Prompt user for settings
Write-Host "Leave an option blank to use the previous value."
$TruncatedPrompt = $(if ($Settings.Prompt.length -gt 50) { $Settings.Prompt.substring(0, 50) + "..." } else { $Settings.Prompt })
$Prompt = Read-Host -Prompt "Generation prompt [$TruncatedPrompt]"
if ($Prompt -eq $null -or $Prompt -eq '') {
    $Prompt = $Settings.Prompt
}

$Steps = Read-Host -Prompt "Target steps [$($Settings.Steps)]"
if ($Steps -eq $null -or $Steps -eq '') {
    $Steps = $Settings.Steps
}

$Iters = Read-Host -Prompt "Number of iterations [$($Settings.Iters)]"
if ($Iters -eq $null -or $Iters -eq '') {
    $Iters = $Settings.Iters
}

$Samples = Read-Host -Prompt "Samples per iteration [$($Settings.Samples)]"
if ($Samples -eq $null -or $Samples -eq '') {
    $Samples = $Settings.Samples
}

$ProgressiveSteps = ,$Steps
for ($i = 0; $i -lt 2; $i++) {
  $NewSteps = [Math]::Ceiling($ProgressiveSteps[0] / 2)
  $ProgressiveSteps = @($NewSteps) + $ProgressiveSteps
}

$Seed = Read-Host -Prompt "Seed (input ? for random) [$($Settings.Seed)]"
if ($Seed -eq $null -or $Seed -eq '') {
    $Seed = $Settings.Seed
}
if ( $Seed -eq '?' ) {
  $Seed = Get-Random
}

Write-Host "Progressively increase steps until target?"
Write-Host "`t* Runs faster generations with fewer steps first, so you can quickly judge your prompt ($($ProgressiveSteps -join ' >> '))"
$ProgressiveGeneration = Read-Host -Prompt "`t(y/n) [$($Settings.ProgressiveGeneration)]"
if ($ProgressiveGeneration -eq $null -or $ProgressiveGeneration -eq '') {
    $ProgressiveGeneration = $Settings.ProgressiveGeneration
}

Write-Host "Use VRAM-optimized script?"
Write-Host "`t* Much slower generation, but uses substantially less VRAM"
$UseOptimizedScript = Read-Host -Prompt "`t(y/n) [$($Settings.UseOptimizedScript)]"
if ($UseOptimizedScript -eq $null -or $UseOptimizedScript -eq '') {
    $UseOptimizedScript = $Settings.UseOptimizedScript
}

# Write the new settings to the settings file
$Settings = @{
    Prompt = $Prompt
    Steps = $Steps
    Iters = $Iters
    Samples = $Samples
    Seed = $Seed
    ProgressiveGeneration = $ProgressiveGeneration
    UseOptimizedScript = $UseOptimizedScript
}
$Settings | ConvertTo-Json -Depth 4 | Set-Content $SettingsPath

if ($ProgressiveGeneration -eq 'y') {
  $Steps = @($ProgressiveSteps)
} else {
  $Steps = @($Steps)
}

$ScriptPath = If ($Settings.UseOptimizedScript -eq 'y') {
    './optimizedSD/optimized_txt2img.py'
} else {
    './scripts/txt2img.py'
}

Write-Host "`r`nRunning with the following settings:"
Write-Host "Prompt: $Prompt"
Write-Host "Steps: $Steps"
Write-Host "Iters: $Iters"
Write-Host "Samples: $Samples"
Write-Host "Seed: $Seed"
Write-Host "Progressive generation: $ProgressiveGeneration`r`n"
Write-Host "$(if ($UseOptimizedScript -eq 'y') { 'Using VRAM-optimized script' } else { 'Using standard script' }) ($ScriptPath)"
Write-Host "`r`n"
Write-Host 'Press Ctrl+C to interrupt.'
Write-Host "`r`n"

$LogPath = './prompt_log.txt'
$Log = If (Test-Path $LogPath) {
    Get-Content $LogPath -Raw
} else {
    ''
}
$Log = $Log + "`r`n$((Get-Date).ToString('yyyy-MM-dd HH:mm:ss')) - Seed: $Seed - Prompt: $Prompt"
Set-Content $LogPath $Log

foreach ($Step in $Steps) {
  python scripts/txt2img.py --prompt $Prompt --H 512 --W 512 --seed $Seed --n_iter $Iters --n_samples $Samples --ddim_steps $Step
}