Update app.py

app.py

@@ -5,6 +5,7 @@ model_1 = gr.load("models/pimpilikipilapi1/NSFW_master")
 model_2 = gr.load("models/DiegoJR1973/NSFW-TrioHMH-Flux")
 model_3 = gr.load("models/prashanth970/flux-lora-uncensored")
 
+
 default_negative_prompt = (
     "Extra limbs, Extra fingers or toes, Disfigured face, Distorted hands, Mutated body parts, "
     "Missing limbs, Asymmetrical features, Blurry face, Poor anatomy, Incorrect proportions, Crooked eyes, "
@@ -18,42 +19,33 @@ default_negative_prompt = (
 )
 
 
-def process_image(model, prompt, negative_prompt=None, supports_negative_prompt=True):
+def generate_image_model_1(prompt, negative_prompt):
+    prompt += " 2k"
     try:
-        if "master" in str(model):
-            prompt += " 2k"
-        elif "Diego" in str(model):
-            prompt += " 8k"
-        else:
-            prompt += " 10k"
-
-
-        if supports_negative_prompt and negative_prompt:
-            return model(prompt, negative_prompt=negative_prompt)
-        else:
-            return model(prompt)
+        return model_1(prompt, negative_prompt=negative_prompt)
     except TypeError:
-        return model(prompt)
-
-
-def generate_images(prompt, negative_prompt):
+        return model_1(prompt)
 
-    output_1 = process_image(model_1, prompt, negative_prompt, True)
-    output_2 = process_image(model_2, prompt, negative_prompt, False)
-    output_3 = process_image(model_3, prompt, negative_prompt, True)
-
-    outputs = []
-    for output in [output_1, output_2, output_3]:
-        if isinstance(output, tuple):
-            outputs.append(None)
-        else:
-            outputs.append(output)
+def generate_image_model_2(prompt, negative_prompt):
+    prompt += " 8k"
+    try:
+        return model_2(prompt, negative_prompt=negative_prompt)
+    except TypeError:
+        return model_2(prompt)
 
-    return outputs
+def generate_image_model_3(prompt, negative_prompt):
+    prompt += " 10k"
+    try:
+        return model_3(prompt, negative_prompt=negative_prompt)
+    except TypeError:
+        return model_3(prompt)
 
-# Gradio Interface
 interface = gr.Interface(
-    fn=generate_images,
+    fn=lambda prompt, negative_prompt: (
+        generate_image_model_1(prompt, negative_prompt),
+        generate_image_model_2(prompt, negative_prompt),
+        generate_image_model_3(prompt, negative_prompt)
+    ),
     inputs=[
         gr.Textbox(label="Type your prompt here: ✍️", placeholder="Describe what you want..."),
         gr.Textbox(label="Negative prompt:", value=default_negative_prompt),
@@ -68,4 +60,4 @@ interface = gr.Interface(
     description="⚠️ Sorry for the inconvenience. The model is currently running on the CPU, which might affect performance. We appreciate your understanding.",
 )
 
-interface.launch(share=True)
+interface.launch()