Update app.py

app.py

@@ -5,10 +5,12 @@ from diffusers import DiffusionPipeline
 from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 from functools import lru_cache
 from PIL import Image
+
 from torchvision import transforms
+from transformers import CLIPFeatureExtractor  # Added missing import
+
 
 
-# Cache pipeline loading to improve performance
 @lru_cache(maxsize=1)
 def load_pipeline():
     # Load base model
@@ -22,7 +24,7 @@ def load_pipeline():
     lora_repo = "strangerzonehf/Flux-Super-Realism-LoRA"
     pipe.load_lora_weights(lora_repo)
     
-    # Load safety checker
+    # Load safety checker and feature extractor
     safety_checker = StableDiffusionSafetyChecker.from_pretrained(
         "CompVis/stable-diffusion-safety-checker"
     )
@@ -30,9 +32,11 @@ def load_pipeline():
         "openai/clip-vit-base-patch32"
     )
     
-    # Optimizations
-    pipe.enable_xformers_memory_efficient_attention()
-    pipe = pipe.to("cuda")
+    # Optimizations: enable memory efficient attention if using GPU
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    if device.type == "cuda":
+        pipe.enable_xformers_memory_efficient_attention()
+    pipe = pipe.to(device)
     
     return pipe, safety_checker, feature_extractor
 
@@ -49,19 +53,19 @@ def generate_image(
 ):
     try:
         progress(0, desc="Initializing...")
-        generator = torch.Generator(device="cuda").manual_seed(seed)
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        generator = torch.Generator(device=device).manual_seed(seed)
         
-        # Auto-add trigger words
+        # Auto-add trigger words if not present
         if "super realism" not in prompt.lower():
             prompt = f"Super Realism, {prompt}"
         
-        # Create callback for progress updates
-        def update_progress(step, _, __):
+        # Define the callback function with the proper signature
+        def update_progress(step, timestep, latents):
             progress((step + 1) / steps, desc="Generating image...")
         
-        # Generate image
         with torch.inference_mode():
-            image = pipe(
+            result = pipe(
                 prompt=prompt,
                 width=width,
                 height=height,
@@ -69,18 +73,21 @@ def generate_image(
                 num_inference_steps=steps,
                 generator=generator,
                 callback=update_progress
-            ).images[0]
+            )
+            image = result.images[0]
         
-        # Safety check
         progress(1, desc="Safety checking...")
+        # Preprocess image for safety checking
         safety_input = feature_extractor(image, return_tensors="pt")
         np_image = np.array(image)
-        safety_result = safety_checker(
+        
+        # Unpack safety checker results (the safety checker returns a tuple)
+        _, nsfw_detected = safety_checker(
             images=[np_image], 
             clip_input=safety_input.pixel_values
         )
         
-        if safety_result.nsfw[0]:
+        if nsfw_detected[0]:
             return Image.new("RGB", (512, 512)), "NSFW content detected"
         
         return image, "Generation successful"
@@ -88,18 +95,17 @@ def generate_image(
     except Exception as e:
         return Image.new("RGB", (512, 512)), f"Error: {str(e)}"
 
-# Create Gradio interface with rate limiting
 with gr.Blocks() as app:
     gr.Markdown("# Flux Super Realism Generator")
     
     with gr.Row():
         with gr.Column():
-            prompt = gr.Textbox(label="Prompt", value="A portrait of a person")
-            seed = gr.Slider(0, 1000, value=42, label="Seed")
-            width = gr.Slider(512, 2048, value=1024, label="Width")
-            height = gr.Slider(512, 2048, value=1024, label="Height")
-            guidance = gr.Slider(1, 20, value=6, label="Guidance Scale")
-            steps = gr.Slider(10, 100, value=28, label="Steps")
+            prompt_input = gr.Textbox(label="Prompt", value="A portrait of a person")
+            seed_input = gr.Slider(0, 1000, value=42, label="Seed")
+            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")
         
         with gr.Column():
@@ -108,13 +114,13 @@ with gr.Blocks() as app:
     
     submit.click(
         generate_image,
-        inputs=[prompt, seed, width, height, guidance, steps],
+        inputs=[prompt_input, seed_input, width_input, height_input, guidance_input, steps_input],
         outputs=[output_image, status]
     )
     
-    # Rate limiting example (1 request every 30 seconds)
+    # Rate limiting: 1 request at a time, with a max queue size of 3
     app.queue(concurrency_count=1, max_size=3).launch()
 
-# For multiple GPU support (advanced)
-# pipe.enable_model_cpu_offload()
-# pipe.enable_sequential_cpu_offload()
+# Uncomment the lines below for advanced multiple GPU support
+pipe.enable_model_cpu_offload()
+pipe.enable_sequential_cpu_offload()