chore: main-backup

README.md

@@ -1,12 +1,83 @@
 ---
-title: Msgxai Hg Api
-emoji: 🦀
+title: Msgxai Hugging Face Inference API
+emoji: 🖼️
 colorFrom: yellow
 colorTo: indigo
 sdk: docker
 pinned: false
 license: mit
-short_description: msgxai backend api
+short_description: Stable Diffusion XL image generation API
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Msgxai Hugging Face Inference API
+
+A custom Hugging Face Inference Endpoint for Stable Diffusion XL image generation.
+
+## Configuration
+
+The API is configured through the `app.conf` JSON file with the following parameters:
+
+```json
+{
+    "model_id": "model-repo-id",       // The Hugging Face model repository ID
+    "name": "your-model-name",         // A name for your model (optional)
+    "prompt": "{prompt}",              // Prompt template with {prompt} placeholder
+    "negative_prompt": "...",          // Default negative prompt
+    "width": 1024,                     // Default image width
+    "height": 768,                     // Default image height
+    "inference_steps": 30,             // Default number of inference steps
+    "guidance_scale": 7,               // Default guidance scale
+    "use_safetensors": true,           // Whether to use safetensors
+    "clip_skip": 0                     // Optional CLIP skip value (0 = disabled)
+}
+```
+
+## API Usage
+
+### Health Check
+```
+GET /
+```
+Returns: `{"status": "healthy"}`
+
+### Generate Image
+```
+POST /predict
+```
+
+Request Body:
+```json
+{
+    "prompt": "your image prompt here",
+    "negative_prompt": "optional negative prompt",
+    "width": 1024,
+    "height": 768,
+    "inference_steps": 30,
+    "guidance_scale": 7,
+    "seed": 42
+}
+```
+
+Response:
+```json
+{
+    "image_base64": "base64-encoded-image-data",
+    "seed": 42
+}
+```
+
+Note: All parameters except `prompt` are optional and will use defaults from `app.conf` if not provided.
+
+## Deployment
+
+1. Configure your `app.conf` file with desired model and parameters
+2. Ensure all dependencies are in `requirements.txt`
+3. Deploy to Hugging Face Inference Endpoints
+
+## Content Filtering
+
+The API includes built-in filtering for child-related content in prompts.
+
+## Environment Variables
+
+- `USE_TORCH_COMPILE`: Set to "1" to enable torch compilation (default: "0")

app.conf

@@ -0,0 +1,11 @@
+{
+    "model_id": "John6666/wai-ani-hentai-pony-v3-sdxl",
+    "name": "hentai-waianiv6-card",
+    "prompt": "score_9, score_8_up, score_7_up,rating_explicit,BREAK, {prompt}",
+    "negative_prompt": "source_furry, source_pony, source_cartoon,3d, blurry, incest, beastiality, children, loli, child, kids, teens, text, logo, timestamp, artist name, artist logo, watermark, web address, copyright name, copyright notice, emblem, comic, title, logo, character name, border, patreon username, signature, webpage, company name, caption, labels, comments",
+    "width": 1024,
+    "height": 768,
+    "inference_steps": 30,
+    "guidance_scale": 7,
+    "use_safetensors": true
+} 

app.py

@@ -1,7 +1,32 @@
-from fastapi import FastAPI
+from fastapi import FastAPI, Request
+from src.handler import EndpointHandler
+import json
 
 app = FastAPI()
 
+# Initialize the handler
+handler = EndpointHandler()
+
 @app.get("/")
 def greet_json():
-    return {"Hello": "World!"}
+    """Simple health check endpoint."""
+    return {"status": "healthy"}
+
+@app.post("/predict")
+async def predict(request: Request):
+    """
+    Main prediction endpoint that processes image generation requests.
+    
+    Args:
+        request (Request): The FastAPI request object
+        
+    Returns:
+        dict: The generated image as base64 and other metadata
+    """
+    # Parse the request data
+    data = await request.json()
+    
+    # Process the request using our handler
+    result = handler(data)
+    
+    return result

requirements.txt

@@ -1,2 +1,9 @@
 fastapi
 uvicorn[standard]
+diffusers
+transformers
+torch
+accelerate
+huggingface_hub
+pillow
+safetensors

src/handler.py

@@ -0,0 +1,194 @@
+import os
+import json
+import random
+import re
+import base64
+from io import BytesIO
+
+import torch
+from huggingface_hub import snapshot_download
+from diffusers import (
+    AutoencoderKL,
+    StableDiffusionXLPipeline,
+    EulerAncestralDiscreteScheduler,
+    DPMSolverSDEScheduler
+)
+from diffusers.models.attention_processor import AttnProcessor2_0
+from PIL import Image
+
+# Global constants
+MAX_SEED = 12211231  # Maximum seed value for random generator
+NUM_IMAGES_PER_PROMPT = 1  # Number of images to generate per prompt
+USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"  # Flag to enable torch compilation
+
+# --- Child-Content Filtering Functions ---
+child_related_regex = re.compile(
+    r'(child|children|kid|kids|baby|babies|toddler|infant|juvenile|minor|underage|preteen|adolescent|youngster|youth|son|daughter|young|kindergarten|preschool|'
+    r'([1-9]|1[0-7])[\s_\-|\.\,]*year(s)?[\s_\-|\.\,]*old|'
+    r'little|small|tiny|short|young|new[\s_\-|\.\,]*born[\s_\-|\.\,]*(boy|girl|male|man|bro|brother|sis|sister))',
+    re.IGNORECASE
+)
+
+def remove_child_related_content(prompt: str) -> str:
+    """Remove any child-related references from the prompt."""
+    # Filter out child-related words/phrases using regex
+    cleaned_prompt = re.sub(child_related_regex, '', prompt)
+    return cleaned_prompt.strip()
+
+def contains_child_related_content(prompt: str) -> bool:
+    """Check if the prompt contains child-related content."""
+    # Use regex to determine if prompt has child-related terms
+    return bool(child_related_regex.search(prompt))
+
+# --- Utility Function: Convert PIL Image to Base64 ---
+def pil_image_to_base64(img: Image.Image) -> str:
+    """Convert a PIL Image to base64 encoded string."""
+    # Create a BytesIO buffer and save the image to it
+    buffered = BytesIO()
+    img.convert("RGB").save(buffered, format="WEBP", quality=90)
+    # Convert buffer to base64 string
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+class EndpointHandler:
+    """
+    Custom handler for Hugging Face Inference Endpoints.
+    This class follows the HF Inference Endpoints specification.
+    """
+    
+    def __init__(self, path="", config=None):
+        """
+        Initialize the handler with model path and configurations.
+        
+        Args:
+            path (str): Path to the model. Not used for this implementation.
+            config (dict, optional): Configuration for the handler. Not used for this implementation.
+        """
+        # Load configuration from app.conf
+        try:
+            with open("app.conf", "r") as f:
+                self.cfg = json.load(f)
+        except Exception as e:
+            print(f"Error loading configuration: {e}")
+            self.cfg = {}
+            
+        # Load the model pipeline
+        print("Loading the model pipeline...")
+        self.pipe = self._load_pipeline_and_scheduler()
+        print("Model loaded successfully!")
+    
+    def _load_pipeline_and_scheduler(self):
+        """Load the Stable Diffusion pipeline and scheduler."""
+        # Get clip_skip from configuration, default to 0
+        clip_skip = self.cfg.get("clip_skip", 0)
+        
+        # Download model files from Hugging Face Hub
+        ckpt_dir = snapshot_download(repo_id=self.cfg["model_id"])
+        
+        # Load the VAE model (for decoding latents)
+        vae = AutoencoderKL.from_pretrained(os.path.join(ckpt_dir, "vae"), torch_dtype=torch.float16)
+        
+        # Load the Stable Diffusion XL pipeline
+        pipe = StableDiffusionXLPipeline.from_pretrained(
+            ckpt_dir,
+            vae=vae,
+            torch_dtype=torch.float16,
+            use_safetensors=self.cfg.get("use_safetensors", True),
+            variant="fp16"
+        )
+        # Move model to GPU
+        pipe = pipe.to("cuda")
+        # Use efficient attention processor
+        pipe.unet.set_attn_processor(AttnProcessor2_0())
+        
+        # Set up samplers/schedulers based on configuration
+        samplers = {
+            "Euler a": EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config),
+            "DPM++ SDE Karras": DPMSolverSDEScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)
+        }
+        # Default to "DPM++ SDE Karras" if not specified
+        pipe.scheduler = samplers.get(self.cfg.get("sampler", "DPM++ SDE Karras"))
+        
+        # Adjust the text encoder layers if needed using clip_skip
+        if clip_skip > 0:
+            pipe.text_encoder.config.num_hidden_layers -= (clip_skip - 1)
+        
+        # Compile model if environment variable is set
+        if USE_TORCH_COMPILE:
+            pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
+            print("Model Compiled!")
+        
+        return pipe
+    
+    def __call__(self, data):
+        """
+        Process the inference request.
+        This is called for each inference request.
+        
+        Args:
+            data: The input data for the inference request
+        
+        Returns:
+            dict: The result of the inference
+        """
+        # Validate that the model is loaded
+        if not self.pipe:
+            return {"error": "Model not loaded. Please check initialization logs."}
+        
+        # Parse the request payload
+        try:
+            if isinstance(data, dict):
+                payload = data
+            else:
+                # Assuming the request is a JSON string
+                payload = json.loads(data)
+        except Exception as e:
+            return {"error": f"Failed to parse request data: {str(e)}"}
+        
+        # Get the prompt from the payload
+        prompt_text = payload.get("prompt", "")
+        if not prompt_text:
+            return {"error": "No prompt provided"}
+        
+        # Apply child-content filtering to the prompt
+        if contains_child_related_content(prompt_text):
+            prompt_text = remove_child_related_content(prompt_text)
+        
+        # Replace placeholder in the prompt template from config
+        combined_prompt = self.cfg.get("prompt", "{prompt}").replace("{prompt}", prompt_text)
+        # Use negative_prompt if provided; otherwise, default to config
+        negative_prompt = payload.get("negative_prompt", self.cfg.get("negative_prompt", ""))
+        
+        # Get parameters from config or override with request params
+        width = int(payload.get("width", self.cfg.get("width", 1024)))
+        height = int(payload.get("height", self.cfg.get("height", 768)))
+        inference_steps = int(payload.get("inference_steps", self.cfg.get("inference_steps", 30)))
+        guidance_scale = float(payload.get("guidance_scale", self.cfg.get("guidance_scale", 7)))
+        
+        # Use provided seed or generate a random one
+        seed = int(payload.get("seed", random.randint(0, MAX_SEED)))
+        generator = torch.Generator(self.pipe.device).manual_seed(seed)
+        
+        try:
+            # Generate the image using the pipeline
+            outputs = self.pipe(
+                prompt=combined_prompt,
+                negative_prompt=negative_prompt,
+                width=width,
+                height=height,
+                guidance_scale=guidance_scale,
+                num_inference_steps=inference_steps,
+                generator=generator,
+                num_images_per_prompt=NUM_IMAGES_PER_PROMPT,
+                output_type="pil"
+            )
+            # Convert the first generated image to base64
+            img_base64 = pil_image_to_base64(outputs.images[0])
+            
+            # Return the response
+            return {"image_base64": img_base64, "seed": seed}
+        
+        except Exception as e:
+            # Log the error and return an error response
+            error_message = f"Image generation failed: {str(e)}"
+            print(error_message)
+            return {"error": error_message}