Update app.py

app.py

@@ -5,11 +5,12 @@ from diffusers import DiffusionPipeline
 from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 from functools import lru_cache
 from PIL import Image
-
+from huggingface_hub import login
 from torchvision import transforms
-from transformers import CLIPFeatureExtractor  # Added missing import
-
+from transformers import CLIPImageProcessor  # Updated import
 
+# Initialize with your Hugging Face token
+login(token="YOUR_HF_TOKEN")
 
 @lru_cache(maxsize=1)
 def load_pipeline():
@@ -24,13 +25,11 @@ def load_pipeline():
     lora_repo = "strangerzonehf/Flux-Super-Realism-LoRA"
     pipe.load_lora_weights(lora_repo)
     
-    # Load safety checker and feature extractor
+    # Load safety checker and image processor
     safety_checker = StableDiffusionSafetyChecker.from_pretrained(
         "CompVis/stable-diffusion-safety-checker"
     )
-    feature_extractor = CLIPFeatureExtractor.from_pretrained(
-        "openai/clip-vit-base-patch32"
-    )
+    image_processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
     
     # Optimizations: enable memory efficient attention if using GPU
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -38,9 +37,9 @@ def load_pipeline():
         pipe.enable_xformers_memory_efficient_attention()
     pipe = pipe.to(device)
     
-    return pipe, safety_checker, feature_extractor
+    return pipe, safety_checker, image_processor
 
-pipe, safety_checker, feature_extractor = load_pipeline()
+pipe, safety_checker, image_processor = load_pipeline()
 
 def generate_image(
     prompt,
@@ -77,11 +76,11 @@ def generate_image(
             image = result.images[0]
         
         progress(1, desc="Safety checking...")
-        # Preprocess image for safety checking
-        safety_input = feature_extractor(image, return_tensors="pt")
+        # Preprocess image for safety checking using the updated image processor
+        safety_input = image_processor(image, return_tensors="pt")
         np_image = np.array(image)
         
-        # Unpack safety checker results (the safety checker returns a tuple)
+        # Unpack safety checker results
         _, nsfw_detected = safety_checker(
             images=[np_image], 
             clip_input=safety_input.pixel_values
@@ -120,6 +119,7 @@ with gr.Blocks() as app:
     
     # Rate limiting: 1 request at a time, with a max queue size of 3
     app.queue(max_size=3).launch()
-# Uncomment the lines below for advanced multiple GPU support
-pipe.enable_model_cpu_offload()
-pipe.enable_sequential_cpu_offload()
+
+# Uncomment for advanced multiple GPU support:
+# pipe.enable_model_cpu_offload()
+# pipe.enable_sequential_cpu_offload()