fixing prompt code

S2I/commons/controller.py

@@ -56,7 +56,7 @@ class Sketch2ImageController():
     
     def artwork(self, options, image, prompt, prompt_template, style_name, seed, val_r, faster, model_name, type_flag, prompt_quality):
         self.load_pipeline(zero_options=options)
-        prompt = prompt_template.replace("{prompt}", prompt)
+        # prompt = prompt_template.replace("{prompt}", prompt)
 
         # if type_flag == 'live-sketch':
         #     img = Image.fromarray(np.array(image["composite"])[:, :, -1])
@@ -79,7 +79,7 @@ class Sketch2ImageController():
         noise = torch.randn((1, 4, H // 8, W // 8), device=c_t.device)
 
         with torch.no_grad():
-            output_image = self.pipe.generate(c_t, prompt, prompt_quality, r=val_r, noise_map=noise, half_model=faster, model_name=model_name)
+            output_image = self.pipe.generate(c_t, prompt, prompt_quality, prompt_template, r=val_r, noise_map=noise, half_model=faster, model_name=model_name)
 
         output_pil = F.to_pil_image(output_image[0].cpu() * 0.5 + 0.5)
 

S2I/modules/sketch2image.py

@@ -13,9 +13,10 @@ class Sketch2ImagePipeline(PrimaryModel):
         super().__init__()
         self.timestep = torch.tensor([999], device="cuda").long()
 
-    def generate(self, c_t, prompt=None, prompt_quality=None, prompt_tokens=None, r=1.0, noise_map=None, half_model=None, model_name=None):
+    def generate(self, c_t, prompt=None, prompt_quality=None, prompt_template=None, prompt_tokens=None, r=1.0, noise_map=None, half_model=None, model_name=None):
         self.from_pretrained(model_name=model_name, r=r)
         prompt_enhanced = self.automatic_enhance_prompt(prompt, prompt_quality)
+        prompt_enhanced = prompt_template.replace("{prompt}", prompt_enhanced)
         assert (prompt is None) != (prompt_tokens is None), "Either prompt or prompt_tokens should be provided"
 
         if half_model == 'float16':