Update app.py

app.py

@@ -9,8 +9,17 @@ device=torch.device('cuda')
 
 # Load the model and LoRA weights
 pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
-#pipe.load_lora_weights("Shakker-Labs/FLUX.1-dev-LoRA-Children-Simple-Sketch", weight_name="FLUX-dev-lora-children-simple-sketch.safetensors")
-pipe.load_lora_weights("Datou1111/shou_xin", weight_name="shou_xin.safetensors")
+print(prompt)
+    
+if style='shou_xin':
+    prompt= "shou_xin, " + prompt
+    pipe.load_lora_weights("Datou1111/shou_xin", weight_name="shou_xin.safetensors")
+else:
+    prompt= "sketched style, " + prompt
+    pipe.load_lora_weights("Shakker-Labs/FLUX.1-dev-LoRA-Children-Simple-Sketch", weight_name="FLUX-dev-lora-children-simple-sketch.safetensors")
+    
+
+
 pipe.fuse_lora(lora_scale=1.5)
 pipe.to("cuda")
 
@@ -22,7 +31,7 @@ NSFW_THRESHOLD = 0.3
 
 # Define the function to generate the sketch
 @spaces.GPU
-def generate_sketch(prompt, num_inference_steps, guidance_scale):
+def generate_sketch(prompt,style, num_inference_steps, guidance_scale):
     # Classify the text for NSFW content
     #text_classification = text_classifier(prompt)
     #print(text_classification)
@@ -31,9 +40,8 @@ def generate_sketch(prompt, num_inference_steps, guidance_scale):
     #for result in text_classification:
     #    if result['label'] == 'nsfw' and result['score'] > NSFW_THRESHOLD: 
     #        return gr.update(visible=False),gr.Text(value="Inappropriate prompt detected. Please try another prompt.")
-    print(prompt)
-    prompt= "sketched style, " + prompt
-    prompt= "shou_xin, " + prompt
+   
+        
     image = pipe("sketched style, " + prompt, 
                  num_inference_steps=num_inference_steps, 
                  guidance_scale=guidance_scale,
@@ -60,6 +68,9 @@ interface = gr.Interface(
     fn=generate_sketch,
     inputs=[
         "text",  # Prompt input
+        gr.Dropdown(
+        ["sketched", "shou_xin"], label="Style"
+        ),
         gr.Slider(5, 50, value=24, step=1, label="Number of Inference Steps"),  # Slider for num_inference_steps
         gr.Slider(1.0, 10.0, value=3.5, step=0.1, label="Guidance Scale")  # Slider for guidance_scale
     ],