From 172995d210f3219a41f665cfea0aa080db5afe2e Mon Sep 17 00:00:00 2001
From: Pedro Cuenca <pedro@huggingface.co>
Date: Wed, 18 Jan 2023 13:28:12 +0100
Subject: [PATCH] Use negative prompts.

---
 Diffusion/Pipeline/Pipeline.swift | 3 ++-
 Diffusion/State.swift             | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/Diffusion/Pipeline/Pipeline.swift b/Diffusion/Pipeline/Pipeline.swift
index caacd9e..570dc9f 100644
--- a/Diffusion/Pipeline/Pipeline.swift
+++ b/Diffusion/Pipeline/Pipeline.swift
@@ -29,12 +29,13 @@ class Pipeline {
         self.pipeline = pipeline
     }
     
-    func generate(prompt: String, scheduler: StableDiffusionScheduler, numInferenceSteps stepCount: Int = 50, seed: UInt32? = nil) throws -> (CGImage, TimeInterval) {
+    func generate(prompt: String, negativePrompt: String = "", scheduler: StableDiffusionScheduler, numInferenceSteps stepCount: Int = 50, seed: UInt32? = nil) throws -> (CGImage, TimeInterval) {
         let beginDate = Date()
         print("Generating...")
         let theSeed = seed ?? UInt32.random(in: 0..<UInt32.max)
         let images = try pipeline.generateImages(
             prompt: prompt,
+            negativePrompt: negativePrompt,
             imageCount: 1,
             stepCount: stepCount,
             seed: theSeed,
diff --git a/Diffusion/State.swift b/Diffusion/State.swift
index fe607e3..3833be2 100644
--- a/Diffusion/State.swift
+++ b/Diffusion/State.swift
@@ -48,7 +48,7 @@ class GenerationContext: ObservableObject {
     func generate() async -> (CGImage, TimeInterval)? {
         guard let pipeline = pipeline else { return nil }
         let seed = self.seed >= 0 ? UInt32(self.seed) : nil
-        return try? pipeline.generate(prompt: positivePrompt, scheduler: scheduler, numInferenceSteps: Int(steps), seed: seed)
+        return try? pipeline.generate(prompt: positivePrompt, negativePrompt: negativePrompt, scheduler: scheduler, numInferenceSteps: Int(steps), seed: seed)
     }
 }