Update app.py

app.py

@@ -5,14 +5,17 @@ import gradio as gr
 from diffusers import StableDiffusionPipeline
 from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 from transformers import CLIPFeatureExtractor
+from huggingface_hub import hf_hub_download
 
 # Configuration
-BASE_MODEL = "hakurei/waifu-diffusion"
+BASE_MODEL = "stabilityai/stable-diffusion-2-1-base"
+LORA_REPO = "Norod78/Flux-LoRA"
+LORA_FILENAME = "flux_lora.safetensors"
 MODEL_CACHE = "model_cache"
 os.makedirs(MODEL_CACHE, exist_ok=True)
 
 def get_pipeline():
-    # Load the safety checker
+    # Load safety checker + feature extractor
     safety_checker = StableDiffusionSafetyChecker.from_pretrained(
         "CompVis/stable-diffusion-safety-checker"
     )
@@ -20,7 +23,7 @@ def get_pipeline():
         "openai/clip-vit-base-patch32"
     )
 
-    # Load the pipeline
+    # Load base pipeline
     pipe = StableDiffusionPipeline.from_pretrained(
         BASE_MODEL,
         torch_dtype=torch.float32,
@@ -30,14 +33,20 @@ def get_pipeline():
         use_safetensors=True
     )
 
-    # Move to CPU and enable memory optimizations
-    pipe = pipe.to("cpu")
+    # Load and apply LoRA weights
+    lora_path = hf_hub_download(
+        repo_id=LORA_REPO,
+        filename=LORA_FILENAME,
+        cache_dir=MODEL_CACHE
+    )
+    pipe.load_lora_weights(lora_path)
+
+    pipe.to("cpu")
     pipe.enable_attention_slicing()
-    pipe.enable_model_cpu_offload()  # ✅ Requires accelerate>=0.17.0
 
     return pipe
 
-# Load once
+# Load model once
 pipeline = get_pipeline()
 
 def generate_image(prompt, negative_prompt="", width=768, height=768, seed=-1, guidance_scale=7.5, num_inference_steps=25):
@@ -47,7 +56,7 @@ def generate_image(prompt, negative_prompt="", width=768, height=768, seed=-1, g
 
     with torch.no_grad():
         output = pipeline(
-            prompt=f"anime style, {prompt}",
+            prompt=f"flux style, {prompt}",  # Stylized prompt
             negative_prompt=negative_prompt,
             width=width,
             height=height,
@@ -60,7 +69,7 @@ def generate_image(prompt, negative_prompt="", width=768, height=768, seed=-1, g
 
 # Gradio UI
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🌀 Anime Image Generator (Safe for Public)")
+    gr.Markdown("# 🌀 Flux-LoRA Anime Image Generator (CPU only)")
     with gr.Row():
         with gr.Column():
             prompt = gr.Textbox(label="Prompt", lines=3)