Update app.py

app.py

@@ -5,23 +5,19 @@ 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 transformers import CLIPImageProcessor  # Updated per deprecation warning
 
-# Initialize with your Hugging Face token
-login(token="YOUR_HF_TOKEN")
+from torchvision import transforms
+from transformers import CLIPImageProcessor  # Updated import
+
+
 
 @lru_cache(maxsize=1)
 def load_pipeline():
-    # Determine device and appropriate torch_dtype
-    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
     base_model = "black-forest-labs/FLUX.1-dev"
     pipe = DiffusionPipeline.from_pretrained(
         base_model,
-        torch_dtype=torch_dtype,
-        low_cpu_mem_usage=True  # Reduce memory usage during load
+        torch_dtype=torch.bfloat16
     )
     
     # Load LoRA weights
@@ -34,19 +30,12 @@ def load_pipeline():
     )
     image_processor = CLIPImageProcessor.from_pretrained("openai/clip-vit-base-patch32")
     
-    # Enable GPU optimizations if on GPU; else, try sequential offloading on CPU
+    # Optimizations: enable memory efficient attention if using GPU
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     if device.type == "cuda":
-        try:
-            pipe.enable_xformers_memory_efficient_attention()
-        except Exception as e:
-            print("Warning: Could not enable xformers memory efficient attention:", e)
-    else:
-        try:
-            pipe.enable_sequential_cpu_offload()
-        except Exception as e:
-            print("Warning: Could not enable sequential CPU offload:", e)
-    
+        pipe.enable_xformers_memory_efficient_attention()
     pipe = pipe.to(device)
+    
     return pipe, safety_checker, image_processor
 
 pipe, safety_checker, image_processor = load_pipeline()
@@ -54,8 +43,8 @@ pipe, safety_checker, image_processor = load_pipeline()
 def generate_image(
     prompt,
     seed=42,
-    width=512,   # Lowered default resolution for reduced memory usage
-    height=512,  # Lowered default resolution for reduced memory usage
+    width=1024,
+    height=1024,
     guidance_scale=6,
     steps=28,
     progress=gr.Progress()
@@ -65,11 +54,14 @@ def generate_image(
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         generator = torch.Generator(device=device).manual_seed(seed)
         
-        # Auto-add the trigger word if not already present
+        # Auto-add trigger words if not present
         if "super realism" not in prompt.lower():
             prompt = f"Super Realism, {prompt}"
         
-        # Generation without a callback since it's not supported
+        # Define the callback function with the proper signature
+        def update_progress(step, timestep, latents):
+            progress((step + 1) / steps, desc="Generating image...")
+        
         with torch.inference_mode():
             result = pipe(
                 prompt=prompt,
@@ -77,25 +69,25 @@ def generate_image(
                 height=height,
                 guidance_scale=guidance_scale,
                 num_inference_steps=steps,
-                generator=generator
+                generator=generator,
+              
             )
             image = result.images[0]
         
-        progress(0.8, desc="Generation complete. Running safety check...")
-        # Preprocess the image for safety checking
+        progress(1, desc="Safety checking...")
+        # Preprocess image for safety checking using the updated image processor
         safety_input = image_processor(image, return_tensors="pt")
         np_image = np.array(image)
         
-        # Run the safety checker; it returns a tuple where the second element is nsfw flags
+        # Unpack safety checker results
         _, nsfw_detected = safety_checker(
-            images=[np_image],
+            images=[np_image], 
             clip_input=safety_input.pixel_values
         )
         
         if nsfw_detected[0]:
             return Image.new("RGB", (512, 512)), "NSFW content detected"
         
-        progress(1, desc="Generation successful")
         return image, "Generation successful"
     
     except Exception as e:
@@ -108,9 +100,8 @@ with gr.Blocks() as app:
         with gr.Column():
             prompt_input = gr.Textbox(label="Prompt", value="A portrait of a person")
             seed_input = gr.Slider(0, 1000, value=42, label="Seed")
-            # Lower the resolution slider range for less memory-intensive generation
-            width_input = gr.Slider(256, 1024, value=512, label="Width")
-            height_input = gr.Slider(256, 1024, value=512, label="Height")
+            width_input = gr.Slider(512, 2048, value=1024, label="Width")
+            height_input = gr.Slider(512, 2048, value=1024, label="Height")
             guidance_input = gr.Slider(1, 20, value=6, label="Guidance Scale")
             steps_input = gr.Slider(10, 100, value=28, label="Steps")
             submit = gr.Button("Generate")
@@ -125,5 +116,9 @@ with gr.Blocks() as app:
         outputs=[output_image, status]
     )
     
-    # Use queue without GPU-specific parameters
+    # Rate limiting: 1 request at a time, with a max queue size of 3
     app.queue(max_size=3).launch()
+
+# Uncomment for advanced multiple GPU support:
+# pipe.enable_model_cpu_offload()
+# pipe.enable_sequential_cpu_offload()