Update app.py

app.py

@@ -55,34 +55,30 @@ import gradio as gr
 from transformers import pipeline
 import torch
 
-# Check if GPU is available
-device = 0 if torch.cuda.is_available() else -1
+# Maximize CPU usage
+torch.set_num_threads(torch.get_num_threads() * 2)
 
-# Load models using transformers
-model1 = pipeline("text-to-image", model="Jonny001/NSFW_master", device=device)
-model2 = pipeline("text-to-image", model="Jonny001/Alita-v1", device=device)
-model3 = pipeline("text-to-image", model="lexa862/NSFWmodel", device=device)
-model4 = pipeline("text-to-image", model="Keltezaa/flux_pussy_NSFW", device=device)
-model5 = pipeline("text-to-image", model="prashanth970/flux-lora-uncensored", device=device)
+# Load models using Hugging Face pipelines
+model1 = pipeline("text-to-image", model="Jonny001/NSFW_master", device_map="auto")
+model2 = pipeline("text-to-image", model="Jonny001/Alita-v1", device_map="auto")
+model3 = pipeline("text-to-image", model="lexa862/NSFWmodel", device_map="auto")
+model4 = pipeline("text-to-image", model="Keltezaa/flux_pussy_NSFW", device_map="auto")
+model5 = pipeline("text-to-image", model="prashanth970/flux-lora-uncensored", device_map="auto")
 
 # Function to generate images
 def generate_images(text, selected_model):
-    if selected_model == "Model 1 (NSFW Master)":
-        model = model1
-    elif selected_model == "Model 2 (Alita)":
-        model = model2
-    elif selected_model == "Model 3 (Lexa NSFW)":
-        model = model3
-    elif selected_model == "Model 4 (Flux NSFW)":
-        model = model4
-    elif selected_model == "Model 5 (Lora Uncensored)":
-        model = model5
-    else:
-        return "Invalid model selection."
-
-    # Generate images
+    models = {
+        "Model 1 (NSFW Master)": model1,
+        "Model 2 (Alita)": model2,
+        "Model 3 (Lexa NSFW)": model3,
+        "Model 4 (Flux NSFW)": model4,
+        "Model 5 (Lora Uncensored)": model5,
+    }
+    
+    model = models.get(selected_model, model1)
     results = []
-    for i in range(2):
+    
+    for i in range(3):
         modified_text = f"{text} variation {i+1}"
         result = model(modified_text)
         results.append(result)
@@ -103,11 +99,14 @@ interface = gr.Interface(
     outputs=[
         gr.Image(label="Generated Image 1"),
         gr.Image(label="Generated Image 2"),
+        gr.Image(label="Generated Image 3"),
     ],
     theme="Yntec/HaleyCH_Theme_Orange",
-    description="⚠ Sorry for the inconvenience. The models are currently running on the CPU, which might affect performance. We appreciate your understanding.",
+    description="⚠ Models are running on CPU for optimized performance. Your patience is appreciated!",
     cache_examples=False,
 )
 
 # Launch the interface
 interface.launch()
+
+