Update app.py

app.py

@@ -6,18 +6,9 @@ import gradio as gr
 import numpy as np
 import spaces
 import torch
-from diffusers import DiffusionPipeline
+from diffusers import AutoPipelineForText2Image
 from PIL import Image
 
-# Make sure PEFT is installed
-try:
-    import peft
-except ImportError:
-    import subprocess
-    print("Installing PEFT library...")
-    subprocess.check_call(["pip", "install", "peft"])
-    import peft
-
 # Create permanent storage directory
 SAVE_DIR = "saved_images"  # Gradio will handle the persistence
 if not os.path.exists(SAVE_DIR):
@@ -25,17 +16,28 @@ if not os.path.exists(SAVE_DIR):
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 repo_id = "black-forest-labs/FLUX.1-dev"
-adapter_id = "seawolf2357/nsfw-detection"  # Changed to Renoir model
+lora_id = "seawolf2357/nsfw-detection"  # LoRA model
 
-# Initialize pipeline
 print("Loading pipeline...")
-# Use DiffusionPipeline instead of FluxPipeline to ensure proper LoRA compatibility
-pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
-print("Loading LoRA weights...")
-# Add low_cpu_mem_usage=False to avoid PEFT assign=True incompatibility issue
-pipeline.load_lora_weights(adapter_id, low_cpu_mem_usage=False)
+# Use AutoPipelineForText2Image which has better compatibility with LoRA loading
+pipeline = AutoPipelineForText2Image.from_pretrained(
+    repo_id, 
+    torch_dtype=torch.bfloat16,
+    use_safetensors=True
+)
 pipeline = pipeline.to(device)
 
+# Try to load the LoRA with direct method (simpler approach)
+print("Loading LoRA weights...")
+try:
+    pipeline.load_lora_weights(lora_id)
+    print("LoRA weights loaded successfully!")
+    lora_loaded = True
+except Exception as e:
+    print(f"Could not load LoRA weights using standard method: {e}")
+    print("Continuing without LoRA functionality.")
+    lora_loaded = False
+
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
@@ -56,28 +58,6 @@ def save_generated_image(image, prompt):
     
     return filepath
 
-def load_generated_images():
-    if not os.path.exists(SAVE_DIR):
-        return []
-    
-    # Load all images from the directory
-    image_files = [os.path.join(SAVE_DIR, f) for f in os.listdir(SAVE_DIR) 
-                  if f.endswith(('.png', '.jpg', '.jpeg', '.webp'))]
-    # Sort by creation time (newest first)
-    image_files.sort(key=lambda x: os.path.getctime(x), reverse=True)
-    return image_files
-
-def load_predefined_images():
-    predefined_images = [
-        "assets/r1.webp",
-        "assets/r2.webp",
-        "assets/r3.webp",
-        "assets/r4.webp",
-        "assets/r5.webp",
-        "assets/r6.webp",
-    ]
-    return predefined_images
-
 # Function to ensure "nsfw" and "[trigger]" are in the prompt
 def process_prompt(prompt):
     # Add "nsfw" prefix if not already present
@@ -112,21 +92,31 @@ def inference(
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device=device).manual_seed(seed)
     
-    # Use joint_attention_kwargs to control LoRA scale 
-    # (FluxPipeline may use a different parameter name but attempt both)
     try:
-        image = pipeline(
-            prompt=processed_prompt,
-            guidance_scale=guidance_scale,
-            num_inference_steps=num_inference_steps,
-            width=width,
-            height=height,
-            generator=generator,
-            joint_attention_kwargs={"scale": lora_scale},
-        ).images[0]
+        # Try with cross_attention_kwargs if LoRA was loaded successfully
+        if lora_loaded:
+            image = pipeline(
+                prompt=processed_prompt,
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps,
+                width=width,
+                height=height,
+                generator=generator,
+                cross_attention_kwargs={"scale": lora_scale}
+            ).images[0]
+        else:
+            # Fall back to standard generation if LoRA wasn't loaded
+            image = pipeline(
+                prompt=processed_prompt,
+                guidance_scale=guidance_scale,
+                num_inference_steps=num_inference_steps,
+                width=width,
+                height=height,
+                generator=generator,
+            ).images[0]
     except Exception as e:
-        # If the above fails, try with cross_attention_kwargs which is more common
-        print(f"First attempt failed with: {e}, trying alternative method...")
+        print(f"Error during inference with cross_attention_kwargs: {e}")
+        # Fall back to standard generation without LoRA parameters
         image = pipeline(
             prompt=processed_prompt,
             guidance_scale=guidance_scale,
@@ -134,17 +124,17 @@ def inference(
             width=width,
             height=height,
             generator=generator,
-            cross_attention_kwargs={"scale": lora_scale},
         ).images[0]
     
     # Save the generated image
     filepath = save_generated_image(image, processed_prompt)
     
-    # Return the image, seed, and updated gallery
-    return image, seed, processed_prompt, load_generated_images()
-
-examples = "A young couple, their bodies glistening with sweat, make love in the rain, the woman"
+    # Return the image, seed, and processed prompt
+    return image, seed, processed_prompt
 
+examples = [
+    "A young couple, their bodies glistening with sweat, make love in the rain, the woman"
+]
 
 # Brighter custom CSS with vibrant colors
 custom_css = """
@@ -187,125 +177,80 @@ button:hover {
     transform: translateY(-2px);
     box-shadow: 0 5px 15px rgba(0,0,0,0.1);
 }
-.tabs {
-    margin-top: 20px;
-}
-.gallery {
-    background-color: rgba(255, 255, 255, 0.5);
-    border-radius: 10px;
-    padding: 10px;
-}
 """
 
 with gr.Blocks(css=custom_css, analytics_enabled=False) as demo:
     gr.HTML('<div class="title">NSFW Detection STUDIO</div>')
     
-    # Model description with the requested content
-
-    with gr.Tabs(elem_classes="tabs") as tabs:
-        with gr.Tab("Generation"):
-            with gr.Column(elem_id="col-container"):
-                with gr.Row():
-                    prompt = gr.Text(
-                        label="Prompt",
-                        show_label=False,
-                        max_lines=1,
-                        placeholder="Enter your prompt (nsfw and [trigger] will be added automatically)",
-                        container=False,
-                    )
-                    run_button = gr.Button("Generate", variant="primary", scale=0)
-
-                result = gr.Image(label="Result", show_label=False)
-                processed_prompt_display = gr.Textbox(label="Processed Prompt", show_label=True)
-
-                with gr.Accordion("Advanced Settings", open=False):
-                    seed = gr.Slider(
-                        label="Seed",
-                        minimum=0,
-                        maximum=MAX_SEED,
-                        step=1,
-                        value=42,
-                    )
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-                    with gr.Row():
-                        width = gr.Slider(
-                            label="Width",
-                            minimum=256,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=32,
-                            value=1024,
-                        )
-                        height = gr.Slider(
-                            label="Height",
-                            minimum=256,
-                            maximum=MAX_IMAGE_SIZE,
-                            step=32,
-                            value=768,
-                        )
+    # Main generation interface
+    with gr.Column(elem_id="col-container"):
+        with gr.Row():
+            prompt = gr.Text(
+                label="Prompt",
+                show_label=False,
+                max_lines=1,
+                placeholder="Enter your prompt (nsfw and [trigger] will be added automatically)",
+                container=False,
+            )
+            run_button = gr.Button("Generate", variant="primary", scale=0)
 
-                    with gr.Row():
-                        guidance_scale = gr.Slider(
-                            label="Guidance scale",
-                            minimum=0.0,
-                            maximum=10.0,
-                            step=0.1,
-                            value=3.5,
-                        )
-                        num_inference_steps = gr.Slider(
-                            label="Number of inference steps",
-                            minimum=1,
-                            maximum=50,
-                            step=1,
-                            value=30,
-                        )
-                        lora_scale = gr.Slider(
-                            label="LoRA scale",
-                            minimum=0.0,
-                            maximum=1.0,
-                            step=0.1,
-                            value=1.0,
-                        )
+        result = gr.Image(label="Result", show_label=False)
+        processed_prompt_display = gr.Textbox(label="Processed Prompt", show_label=True)
 
-                gr.Examples(
-                    examples=examples,
-                    inputs=[prompt],
-                    outputs=[result, seed, processed_prompt_display],
-                )
-
-        with gr.Tab("Gallery"):
-            gallery_header = gr.Markdown("### Your Generated Images")
-            generated_gallery = gr.Gallery(
-                label="Generated Images",
-                columns=3,
-                show_label=False,
-                value=load_generated_images(),
-                elem_id="generated_gallery",
-                elem_classes="gallery",
-                height="auto"
+        with gr.Accordion("Advanced Settings", open=False):
+            seed = gr.Slider(
+                label="Seed",
+                minimum=0,
+                maximum=MAX_SEED,
+                step=1,
+                value=42,
             )
-            refresh_btn = gr.Button("🔄 Refresh Gallery", variant="primary")
+            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
 
-    # Add sample gallery section at the bottom
-    gr.Markdown("### Pierre-Auguste Renoir Style Examples")
-    predefined_gallery = gr.Gallery(
-        label="Sample Images", 
-        columns=3,
-        rows=2,
-        show_label=False, 
-        value=load_predefined_images(),
-        elem_classes="gallery"
-    )
+            with gr.Row():
+                width = gr.Slider(
+                    label="Width",
+                    minimum=256,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=1024,
+                )
+                height = gr.Slider(
+                    label="Height",
+                    minimum=256,
+                    maximum=MAX_IMAGE_SIZE,
+                    step=32,
+                    value=768,
+                )
 
-    # Event handlers
-    def refresh_gallery():
-        return load_generated_images()
+            with gr.Row():
+                guidance_scale = gr.Slider(
+                    label="Guidance scale",
+                    minimum=0.0,
+                    maximum=10.0,
+                    step=0.1,
+                    value=3.5,
+                )
+                num_inference_steps = gr.Slider(
+                    label="Number of inference steps",
+                    minimum=1,
+                    maximum=50,
+                    step=1,
+                    value=30,
+                )
+                lora_scale = gr.Slider(
+                    label="LoRA scale",
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.1,
+                    value=1.0,
+                )
 
-    refresh_btn.click(
-        fn=refresh_gallery,
-        inputs=None,
-        outputs=generated_gallery,
-    )
+        gr.Examples(
+            examples=examples,
+            inputs=[prompt],
+            outputs=[result, seed, processed_prompt_display],
+        )
 
     gr.on(
         triggers=[run_button.click, prompt.submit],
@@ -320,7 +265,7 @@ with gr.Blocks(css=custom_css, analytics_enabled=False) as demo:
             num_inference_steps,
             lora_scale,
         ],
-        outputs=[result, seed, processed_prompt_display, generated_gallery],
+        outputs=[result, seed, processed_prompt_display],
     )
 
 demo.queue()