Update app.py

app.py

@@ -7,16 +7,21 @@ from functools import lru_cache
 from PIL import Image
 
 from torchvision import transforms
-from transformers import CLIPImageProcessor  # Updated import
+from transformers import CLIPImageProcessor  # Updated per deprecation warning
+import os
 
 
 @lru_cache(maxsize=1)
 def load_pipeline():
-    # Load base model
+    # Determine device and set torch_dtype accordingly
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    torch_dtype = torch.bfloat16 if device.type == "cuda" else torch.float32
+
+    # Load base model with the appropriate dtype
     base_model = "black-forest-labs/FLUX.1-dev"
     pipe = DiffusionPipeline.from_pretrained(
         base_model,
-        torch_dtype=torch.bfloat16
+        torch_dtype=torch_dtype
     )
     
     # Load LoRA weights
@@ -29,12 +34,14 @@ def load_pipeline():
     )
     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")
+    # Enable GPU-only optimizations if a GPU is available
     if device.type == "cuda":
-        pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to(device)
+        try:
+            pipe.enable_xformers_memory_efficient_attention()
+        except Exception as e:
+            print("Warning: Could not enable xformers memory efficient attention:", e)
     
+    pipe = pipe.to(device)
     return pipe, safety_checker, image_processor
 
 pipe, safety_checker, image_processor = load_pipeline()
@@ -53,11 +60,11 @@ def generate_image(
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         generator = torch.Generator(device=device).manual_seed(seed)
         
-        # Auto-add trigger words if not present
+        # Ensure the trigger word is present
         if "super realism" not in prompt.lower():
             prompt = f"Super Realism, {prompt}"
         
-        # Define the callback function with the proper signature
+        # Define a callback to update progress
         def update_progress(step, timestep, latents):
             progress((step + 1) / steps, desc="Generating image...")
         
@@ -74,11 +81,11 @@ def generate_image(
             image = result.images[0]
         
         progress(1, desc="Safety checking...")
-        # Preprocess image for safety checking using the updated image processor
+        # Preprocess the image for safety checking
         safety_input = image_processor(image, return_tensors="pt")
         np_image = np.array(image)
         
-        # Unpack safety checker results
+        # Run the safety checker
         _, nsfw_detected = safety_checker(
             images=[np_image], 
             clip_input=safety_input.pixel_values
@@ -115,9 +122,5 @@ with gr.Blocks() as app:
         outputs=[output_image, status]
     )
     
-    # Rate limiting: 1 request at a time, with a max queue size of 3
+    # Queue without GPU-specific arguments
     app.queue(max_size=3).launch()
-
-# Uncomment for advanced multiple GPU support:
-# pipe.enable_model_cpu_offload()
-# pipe.enable_sequential_cpu_offload()