update enhance prompt

S2I/commons/controller.py

@@ -58,16 +58,11 @@ class Sketch2ImageController():
         self.load_pipeline(zero_options=options)
         # prompt = prompt_template.replace("{prompt}", prompt)
 
-        # if type_flag == 'live-sketch':
-        #     img = Image.fromarray(np.array(image["composite"])[:, :, -1])
-        # elif type_flag == 'url-sketch':
-        #     img = image["composite"]
-
-        if type_flag == 'URL':
-            img = image["composite"]
-        else:
+        if type_flag == 'live-sketch':
             img = Image.fromarray(np.array(image["composite"])[:, :, -1])
-
+        elif type_flag == 'url-sketch':
+            img = image["composite"]
+        
         img = img.convert("RGB")
         img = img.resize((512, 512))
 
@@ -83,9 +78,5 @@ class Sketch2ImageController():
 
         output_pil = F.to_pil_image(output_image[0].cpu() * 0.5 + 0.5)
 
-        # if type_flag == 'live-sketch':
-        #     input_uri = self.pil_image_to_data_uri(Image.fromarray(255 - np.array(img)))
-        # else:
-        #     input_uri = self.pil_image_to_data_uri(img)
 
         return output_pil

S2I/modules/models.py

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

S2I/modules/sketch2image.py

@@ -75,8 +75,8 @@ class Sketch2ImagePipeline(PrimaryModel):
         self.global_unet.set_adapters(["default"], weights=[r])
         set_weights_and_activate_adapters(self.global_vae, ["vae_skip"], [r])
 
-    def automatic_enhance_prompt(self, input_prompt, model_choice):
-        if model_choice == "short-sentences":
+    def automatic_enhance_prompt(self, input_prompt, prompt_quality):
+        if prompt_quality:
             result = self.global_medium_prompt("Enhance the description: " + input_prompt)
             enhanced_text = result[0]['summary_text']
             
@@ -87,10 +87,6 @@ class Sketch2ImagePipeline(PrimaryModel):
                 remaining_text = enhanced_text[match.end():].strip()
                 modified_sentence = match.group(1).capitalize()
                 enhanced_text = modified_sentence + ' ' + remaining_text
-        else:
-            result = self.global_long_prompt("Enhance the description: " + input_prompt)
-            enhanced_text = result[0]['summary_text']
-    
         return enhanced_text
 
     def _move_to_cpu(self, module):

app.py

@@ -260,6 +260,7 @@ with gr.Blocks(css=css, theme="NoCrypt/miku@1.2.1") as demo:
                 show_download_button=True,
             )
             with gr.Group():
+                use_enhancer = gr.Checkbox(label="Use Automatic Prompt High-Quality", value=False)
                 prompt = gr.Textbox(label="Personalized Text", value="", show_label=True)
                 with gr.Row():
                     run_button = gr.Button("Generate 🪄", min_width=5, variant='primary')
@@ -267,22 +268,12 @@ with gr.Blocks(css=css, theme="NoCrypt/miku@1.2.1") as demo:
                     clear_button = gr.Button("Reset Sketch Session", min_width=10, variant='primary')
                 with gr.Accordion("S2I Advances Option", open=True):
                         with gr.Row():
-                                # input_type = gr.Radio(
-                                #     choices=["live-sketch", "url-sketch"],
-                                #     value="live-sketch",
-                                #     label="Type Sketch2Image models",
-                                #     interactive=True)
-
-                                input_type = gr.Textbox(
-                                    label="Check URL or Real-time Input",
-                                    interactive=True)
-
-                                prompt_quality = gr.Radio(
-                                    choices=["short-sentences", "long-sentences"],
-                                    value="short-sentences",
-                                    label="Long/Short of Text Prompt",
+                                input_type = gr.Radio(
+                                    choices=["live-sketch", "url-sketch"],
+                                    value="live-sketch",
+                                    label="Type Sketch2Image models",
                                     interactive=True)
-                                    
+                                                                        
                                 style = gr.Dropdown(
                                     label="Style",
                                     choices=controller.STYLE_NAMES,