Update app.py

app.py

@@ -1,48 +1,52 @@
 import os
 import random
-import gradio as gr
 import torch
+import gradio as gr
 from diffusers import StableDiffusionPipeline
+from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
+from transformers import CLIPFeatureExtractor
 
-# Use the anime-specific Waifu Diffusion model
+# Configuration
 BASE_MODEL = "hakurei/waifu-diffusion"
 MODEL_CACHE = "model_cache"
 os.makedirs(MODEL_CACHE, exist_ok=True)
 
 def get_pipeline():
-    # Load the Waifu Diffusion model (anime style) on CPU
+    # Load the safety checker
+    safety_checker = StableDiffusionSafetyChecker.from_pretrained(
+        "CompVis/stable-diffusion-safety-checker"
+    )
+    feature_extractor = CLIPFeatureExtractor.from_pretrained(
+        "openai/clip-vit-base-patch32"
+    )
+
+    # Load the pipeline
     pipe = StableDiffusionPipeline.from_pretrained(
         BASE_MODEL,
         torch_dtype=torch.float32,
         cache_dir=MODEL_CACHE,
-        safety_checker=None,
+        safety_checker=safety_checker,
+        feature_extractor=feature_extractor,
         use_safetensors=True
     )
+
     # Move to CPU and enable memory optimizations
     pipe = pipe.to("cpu")
     pipe.enable_attention_slicing()
-    pipe.enable_model_cpu_offload()
+    pipe.enable_model_cpu_offload()  # ✅ Requires accelerate>=0.17.0
+
     return pipe
 
-# Load pipeline once at start
+# Load once
 pipeline = get_pipeline()
 
-def generate_image(
-    prompt: str,
-    negative_prompt: str = "",
-    width: int = 768,
-    height: int = 768,
-    seed: int = -1,
-    guidance_scale: float = 7.5,
-    num_inference_steps: int = 25
-):
-    # Generate a random seed if none provided
+def generate_image(prompt, negative_prompt="", width=768, height=768, seed=-1, guidance_scale=7.5, num_inference_steps=25):
     if seed == -1:
-        seed = random.randint(0, 2**31-1)
+        seed = random.randint(0, 2**31 - 1)
     generator = torch.Generator(device="cpu").manual_seed(seed)
-    # Generate the image with the pipeline
+
     with torch.no_grad():
-        image = pipeline(
+        output = pipeline(
             prompt=f"anime style, {prompt}",
             negative_prompt=negative_prompt,
             width=width,
@@ -50,26 +54,30 @@ def generate_image(
             guidance_scale=guidance_scale,
             num_inference_steps=num_inference_steps,
             generator=generator
-        ).images[0]
+        )
+        image = output.images[0]
     return image, seed
 
-# Gradio Interface
+# Gradio UI
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🌀 Anime Image Generator (CPU only)")
+    gr.Markdown("# 🌀 Anime Image Generator (Safe for Public)")
     with gr.Row():
         with gr.Column():
-            prompt = gr.Textbox(label="Prompt", lines=3, placeholder="e.g., 1girl, pink hair, night sky...")
-            negative_prompt = gr.Textbox(label="Negative Prompt", value="blurry, lowres, disfigured")
+            prompt = gr.Textbox(label="Prompt", lines=3)
+            negative_prompt = gr.Textbox(label="Negative Prompt", value="blurry, lowres, bad anatomy")
             generate_btn = gr.Button("Generate", variant="primary")
+
             with gr.Accordion("Advanced", open=False):
                 width = gr.Slider(512, 1024, value=768, step=64, label="Width")
                 height = gr.Slider(512, 1024, value=768, step=64, label="Height")
                 guidance = gr.Slider(1.0, 15.0, value=7.5, step=0.5, label="Guidance")
                 steps = gr.Slider(15, 50, value=25, step=1, label="Steps")
                 seed = gr.Number(label="Seed", value=-1)
+
         with gr.Column():
             output_image = gr.Image(label="Result", type="pil")
             used_seed = gr.Textbox(label="Used Seed")
+
     generate_btn.click(
         generate_image,
         inputs=[prompt, negative_prompt, width, height, seed, guidance, steps],
@@ -77,4 +85,4 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch()