Update app.py

app.py

@@ -4,46 +4,43 @@ import random
 import spaces #[uncomment to use ZeroGPU]
 from diffusers import DiffusionPipeline, DPMSolverSDEScheduler
 import torch
-from transformers import AutoModelForObjectDetection, AutoImageProcessor
+from huggingface_hub import hf_hub_download
+from ultralytics import YOLO
+import cv2
+from PIL import Image
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model_repo_id = "John6666/wai-ani-nsfw-ponyxl-v8-sdxl"  # Your diffusion model
 
 # Load your main diffusion pipeline
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
-
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
+pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch.float16)
 pipe.scheduler = DPMSolverSDEScheduler.from_config(pipe.scheduler.config, algorithm_type="dpmsolver++", solver_order=2, use_karras_sigmas=True)
 pipe = pipe.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
-# Load ADetailer model (from Hugging Face)
-adetailer_model_id = "Bingsu/adetailer"
-adetailer_model = AutoModelForObjectDetection.from_pretrained(adetailer_model_id)
-adetailer_processor = AutoImageProcessor.from_pretrained(adetailer_model_id)
+# Download the ADetailer YOLOv8 face detection model
+yolo_model_path = hf_hub_download(repo_id="Bingsu/adetailer", filename="face_yolov8n.pt")
+yolo_model = YOLO(yolo_model_path)
 
 def fix_eyes_with_adetailer(image):
-    # Convert image to format for ADetailer
-    pixel_values = adetailer_processor(images=image, return_tensors="pt").pixel_values
-    pixel_values = pixel_values.to(device)
+    # Convert PIL image to OpenCV format for YOLO
+    img = np.array(image)
+    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
     
-    # Run ADetailer on the image
-    with torch.no_grad():
-        outputs = adetailer_model(pixel_values=pixel_values)
+    # Run the YOLO model on the image
+    results = yolo_model(img)
     
-    # Post-process the outputs and apply the fixes (if any)
-    corrected_image = image  # Placeholder for the actual post-processing
-    # Apply fixes based on the detection and correction model outputs
-    # This step requires actual ADetailer implementation details for correcting eyes.
-
-    return corrected_image  # Return the corrected image
+    # Visualize and process the output
+    pred = results[0].plot()  # Draw bounding boxes and other detections
+    pred = cv2.cvtColor(pred, cv2.COLOR_BGR2RGB)
+    
+    # Convert the processed image back to PIL format
+    corrected_image = Image.fromarray(pred)
+    return corrected_image
 
-@spaces.GPU  #[uncomment to use ZeroGPU]
+@spaces.GPU #[uncomment to use ZeroGPU]
 def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
 
     if randomize_seed:
@@ -51,6 +48,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         
     generator = torch.Generator().manual_seed(seed)
     
+    # Generate the initial image with the diffusion model
     image = pipe(
         prompt=prompt, 
         negative_prompt=negative_prompt,
@@ -61,7 +59,7 @@ def infer(prompt, negative_prompt, seed, randomize_seed, width, height, guidance
         generator=generator
     ).images[0] 
 
-    # Apply ADetailer to fix eyes after generating the image
+    # Apply ADetailer to fix the eyes after generating the image
     corrected_image = fix_eyes_with_adetailer(image)
     
     return corrected_image, seed