Update app.py

app.py

@@ -52,20 +52,21 @@
 # interface.launch()
 
 import gradio as gr
+from transformers import pipeline
 import torch
 
 # Check if GPU is available
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = 0 if torch.cuda.is_available() else -1
 
-# Load models on GPU if available, otherwise fallback to CPU
-model1 = gr.load("models/Jonny001/NSFW_master", device=device)  # GPU
-model2 = gr.load("models/Jonny001/Alita-v1", device=device)     # GPU
-model3 = gr.load("models/lexa862/NSFWmodel", device=device)     # GPU
-model4 = gr.load("models/Keltezaa/flux_pussy_NSFW", device=device)  # GPU
-model5 = gr.load("models/prashanth970/flux-lora-uncensored", device=device)  # GPU
+# 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)
 
+# Function to generate images
 def generate_images(text, selected_model):
-    # Model selection logic
     if selected_model == "Model 1 (NSFW Master)":
         model = model1
     elif selected_model == "Model 2 (Alita)":
@@ -78,14 +79,14 @@ def generate_images(text, selected_model):
         model = model5
     else:
         return "Invalid model selection."
-    
-    # Generate two variations for each input prompt
+
+    # Generate images
     results = []
     for i in range(2):
         modified_text = f"{text} variation {i+1}"
         result = model(modified_text)
         results.append(result)
-    
+
     return results
 
 # Gradio interface
@@ -108,4 +109,5 @@ interface = gr.Interface(
     cache_examples=False,
 )
 
+# Launch the interface
 interface.launch()