Update app.py

app.py

@@ -2,7 +2,7 @@ import cv2
 import torch
 import numpy as np
 import gradio as gr
-from diffusers import StableDiffusionXLPipeline
+from diffusers import StableDiffusionPipeline  # Use SD 2.1 instead of SDXL
 from insightface.app import FaceAnalysis
 from huggingface_hub import hf_hub_download
 import os
@@ -73,15 +73,15 @@ ip_adapter_weights = download_file(
     ip_adapter_path
 )
 
-# Load the pipeline with verbose logging and retry logic
-logger.info("Loading Stable Diffusion XL base model...")
+# Load the pipeline with SD 2.1
+logger.info("Loading Stable Diffusion 2.1 base model...")
 try:
     max_retries = 3
     for attempt in range(max_retries):
         try:
-            logger.info(f"Attempt {attempt + 1}/{max_retries}: Loading SDXL model...")
-            pipe = StableDiffusionXLPipeline.from_pretrained(
-                "stabilityai/stable-diffusion-xl-base-1.0",
+            logger.info(f"Attempt {attempt + 1}/{max_retries}: Loading SD 2.1 model...")
+            pipe = StableDiffusionPipeline.from_pretrained(
+                "stabilityai/stable-diffusion-2-1",
                 torch_dtype=dtype,
                 safety_checker=None,
                 local_files_only=False,
@@ -89,7 +89,7 @@ try:
                 variant="fp16",
                 use_safetensors=True
             )
-            logger.info("SDXL base model loaded successfully.")
+            logger.info("SD 2.1 base model loaded successfully.")
             break
         except Exception as e:
             logger.error(f"Load attempt {attempt + 1} failed: {e}")
@@ -97,9 +97,9 @@ try:
                 logger.info("Retrying in 5 seconds...")
                 time.sleep(5)
             else:
-                raise RuntimeError(f"Failed to load SDXL model after {max_retries} attempts: {e}")
+                raise RuntimeError(f"Failed to load SD 2.1 model after {max_retries} attempts: {e}")
 except Exception as e:
-    logger.error(f"Failed to load SDXL base model: {e}")
+    logger.error(f"Failed to load SD 2.1 base model: {e}")
     raise
 
 # Load IP-Adapter