fixing design block

S2I/commons/controller.py

@@ -56,8 +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_enhanced = self.pipe.automatic_enhance_prompt(prompt, prompt_quality)
-        prompt_enhanced = prompt_template.replace("{prompt}", prompt_enhanced)
+        prompt = prompt_template.replace("{prompt}", prompt)
 
         # if type_flag == 'live-sketch':
         #     img = Image.fromarray(np.array(image["composite"])[:, :, -1])
@@ -80,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_enhanced, r=val_r, noise_map=noise, half_model=faster, model_name=model_name)
+            output_image = self.pipe.generate(c_t, prompt, prompt_quality, 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/models.py

@@ -64,12 +64,11 @@ class PrimaryModel:
             sd = torch.load(p_ckpt, map_location="cpu")
             return sd
     def from_pretrained(self, model_name, r):
-        if model_name is None and r is None:
-            if self.global_medium_prompt is None:
-                self.global_medium_prompt = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device='cuda')
+        if self.global_medium_prompt is None:
+            self.global_medium_prompt = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance", device='cuda')
 
-            if self.global_long_prompt is None:
-                self.global_long_prompt = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device='cuda')
+        if self.global_long_prompt is None:
+            self.global_long_prompt = pipeline("summarization", model="gokaygokay/Lamini-Prompt-Enchance-Long", device='cuda')
 
         if self.global_tokenizer is None:
             self.global_tokenizer = AutoTokenizer.from_pretrained("myn0908/stable-diffusion-3", subfolder="tokenizer_2")

S2I/modules/sketch2image.py

@@ -13,14 +13,15 @@ class Sketch2ImagePipeline(PrimaryModel):
         super().__init__()
         self.timestep = torch.tensor([999], device="cuda").long()
 
-    def generate(self, c_t, prompt=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_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)
         assert (prompt is None) != (prompt_tokens is None), "Either prompt or prompt_tokens should be provided"
 
         if half_model == 'float16':
-            output_image = self._generate_fp16(c_t, prompt, prompt_tokens, r, noise_map)
+            output_image = self._generate_fp16(c_t, prompt_enhanced, prompt_tokens, r, noise_map)
         else:
-            output_image = self._generate_full_precision(c_t, prompt, prompt_tokens, r, noise_map)
+            output_image = self._generate_full_precision(c_t, prompt_enhanced, prompt_tokens, r, noise_map)
 
         return output_image
 
@@ -74,7 +75,6 @@ class Sketch2ImagePipeline(PrimaryModel):
         set_weights_and_activate_adapters(self.global_vae, ["vae_skip"], [r])
 
     def automatic_enhance_prompt(self, input_prompt, model_choice):
-        self.from_pretrained(model_name=None, r=None)
         if model_choice == "short-sentences":
             result = self.global_medium_prompt("Enhance the description: " + input_prompt)
             enhanced_text = result[0]['summary_text']