Update app.py

app.py

@@ -27,12 +27,13 @@ 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
 
-# Initialize pipeline with PEFT support
+# Initialize pipeline
 print("Loading pipeline...")
-# Add 'use_peft=True' parameter to enable PEFT support
-pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16, use_peft=True)
+# Use DiffusionPipeline instead of FluxPipeline to ensure proper LoRA compatibility
+pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch.bfloat16)
 print("Loading LoRA weights...")
-pipeline.load_lora_weights(adapter_id)
+# Add low_cpu_mem_usage=False to avoid PEFT assign=True incompatibility issue
+pipeline.load_lora_weights(adapter_id, low_cpu_mem_usage=False)
 pipeline = pipeline.to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
@@ -111,15 +112,30 @@ def inference(
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device=device).manual_seed(seed)
     
-    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]
+    # 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]
+    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...")
+        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]
     
     # Save the generated image
     filepath = save_generated_image(image, processed_prompt)