Prepare code for Hugging Face Inference Endpoint deployment

.gitattributes

@@ -1,35 +0,0 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

.python-version

@@ -1 +0,0 @@
-msgxai-hg-api

Dockerfile

@@ -1,16 +0,0 @@
-# Read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
-# you will also find guides on how best to write your Dockerfile
-
-FROM python:3.9
-
-RUN useradd -m -u 1000 user
-USER user
-ENV PATH="/home/user/.local/bin:$PATH"
-
-WORKDIR /app
-
-COPY --chown=user ./requirements.txt requirements.txt
-RUN pip install --no-cache-dir --upgrade -r requirements.txt
-
-COPY --chown=user . /app
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,83 +1,125 @@
----
-title: Msgxai Hugging Face Inference API
-emoji: 🖼️
-colorFrom: yellow
-colorTo: indigo
-sdk: docker
-pinned: false
-license: mit
-short_description: Stable Diffusion XL image generation API
----
+# Hugging Face Inference API for Stable Diffusion XL
 
-# Msgxai Hugging Face Inference API
+This repository contains a text-to-image generation API designed to be deployed on Hugging Face Inference Endpoints, using Stable Diffusion XL models for image generation.
 
-A custom Hugging Face Inference Endpoint for Stable Diffusion XL image generation.
+## Features
+
+- Compatible with Hugging Face Inference Endpoints
+- Stable Diffusion XL (SDXL) model for high-quality image generation
+- Content filtering for safe image generation
+- Configurable image dimensions (default: 1024x768)
+- Base64-encoded image output
+- Performance optimizations (torch.compile, attention processors)
+
+## Project Structure
+
+The codebase has been simplified to only use a single file:
+
+- `handler.py`: Contains the `EndpointHandler` class that implements the Hugging Face Inference Endpoints interface. This file also includes a built-in FastAPI server for local development.
 
 ## Configuration
 
-The API is configured through the `app.conf` JSON file with the following parameters:
+The service is configured via 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)
+  "model_id": "your-huggingface-model-id",
+  "prompt": "template with {prompt} placeholder",
+  "negative_prompt": "default negative prompt",
+  "inference_steps": 30,
+  "guidance_scale": 7,
+  "use_safetensors": true,
+  "width": 1024,
+  "height": 768
 }
 ```
 
 ## API Usage
 
-### Health Check
-```
-GET /
-```
-Returns: `{"status": "healthy"}`
+### Hugging Face Inference Endpoints Format
 
-### Generate Image
-```
-POST /predict
-```
+When deployed to Hugging Face Inference Endpoints, the API accepts requests in the following format:
 
-Request Body:
 ```json
 {
-    "prompt": "your image prompt here",
+  "inputs": "your prompt here",
+  "parameters": {
     "negative_prompt": "optional negative prompt",
-    "width": 1024,
-    "height": 768,
+    "seed": 12345,
     "inference_steps": 30,
     "guidance_scale": 7,
-    "seed": 42
+    "width": 1024,
+    "height": 768
+  }
 }
 ```
 
-Response:
+Response format:
+```json
+[
+  {
+    "generated_image": "base64-encoded-image",
+    "seed": 12345
+  }
+]
+```
+
+### Local Development Format
+
+When running locally, you can use the same format as above, or a simplified format:
+
 ```json
 {
-    "image_base64": "base64-encoded-image-data",
-    "seed": 42
+  "prompt": "your prompt here",
+  "negative_prompt": "optional negative prompt",
+  "seed": 12345,
+  "inference_steps": 30,
+  "guidance_scale": 7,
+  "width": 1024,
+  "height": 768
 }
 ```
 
-Note: All parameters except `prompt` are optional and will use defaults from `app.conf` if not provided.
+Response format from the local server:
+```json
+[
+  {
+    "generated_image": "base64-encoded-image",
+    "seed": 12345
+  }
+]
+```
+
+## Deployment on Hugging Face Inference Endpoints
 
-## Deployment
+1. Push this repository to Hugging Face Hub or your Git repository
+2. Create a new Inference Endpoint on Hugging Face
+3. Select this repository as the source
+4. Configure compute resources (recommended: GPU with at least 16GB VRAM)
+5. Deploy the endpoint
 
-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
+### Required Files
 
-## Content Filtering
+For deployment on Hugging Face Inference Endpoints, you need:
+- `handler.py` - Contains the `EndpointHandler` class implementation
+- `requirements.txt` - Lists the Python dependencies
+- `app.conf` - Contains configuration parameters
 
-The API includes built-in filtering for child-related content in prompts.
+Note: A `Procfile` is not needed for Hugging Face Inference Endpoints deployment, as the service automatically detects and uses the `EndpointHandler` class.
+
+## Local Development
+
+1. Install dependencies: `pip install -r requirements.txt`
+2. Run the API locally: `python handler.py [--port PORT] [--host HOST]`
+3. The API will be available at http://localhost:8000
+
+The local server uses the FastAPI implementation included in `handler.py` that provides the same functionality as the Hugging Face Inference Endpoints interface.
 
 ## Environment Variables
 
-- `USE_TORCH_COMPILE`: Set to "1" to enable torch compilation (default: "0")
+- `PORT`: Port to run the server on (default: 8000)
+- `USE_TORCH_COMPILE`: Set to "1" to enable torch.compile for performance (default: "0")
+
+## License
+
+This project is licensed under the terms of the MIT license. 

app.conf

@@ -3,9 +3,9 @@
     "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
-} 
+    "use_safetensors": true,
+    "width": 1024,
+    "height": 768
+}

app.py

@@ -1,32 +0,0 @@
-from fastapi import FastAPI, Request
-from src.handler import EndpointHandler
-import json
-
-app = FastAPI()
-
-# Initialize the handler
-handler = EndpointHandler()
-
-@app.get("/")
-def greet_json():
-    """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

src/handler.py → handler.py

@@ -53,6 +53,10 @@ class EndpointHandler:
     """
     Custom handler for Hugging Face Inference Endpoints.
     This class follows the HF Inference Endpoints specification.
+    
+    For Hugging Face Inference Endpoints, only this class is needed.
+    It provides both the initialization (__init__) and inference (__call__) methods
+    required by the Hugging Face Inference API.
     """
     
     def __init__(self, path="", config=None):
@@ -60,13 +64,20 @@ class EndpointHandler:
         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.
+            path (str): Path to the model directory (used by HF Inference Endpoints).
+            config (dict, optional): Configuration for the handler, passed by HF Inference Endpoints.
         """
-        # Load configuration from app.conf
+        # Load configuration from app.conf or use provided config
         try:
-            with open("app.conf", "r") as f:
-                self.cfg = json.load(f)
+            if config:
+                # Use config provided by HF Inference Endpoints
+                self.cfg = config
+            else:
+                # Try to load from app.conf as fallback
+                config_path = os.path.join(path, "app.conf") if path else "app.conf"
+                with open(config_path, "r") as f:
+                    self.cfg = json.load(f)
+            print("Configuration loaded successfully")
         except Exception as e:
             print(f"Error loading configuration: {e}")
             self.cfg = {}
@@ -122,16 +133,18 @@ class EndpointHandler:
     def __call__(self, data):
         """
         Process the inference request.
-        This is called for each inference request.
+        This is called for each inference request by the Hugging Face Inference API.
         
         Args:
             data: The input data for the inference request
-        
+                 For HF Inference Endpoints, this is typically a dict with "inputs" field
+                 
         Returns:
-            dict: The result of the inference
+            list: A list containing the generated image as base64 string and seed
+                 This follows the HF Inference Endpoints output format
         """
         # Validate that the model is loaded
-        if not self.pipe:
+        if not hasattr(self, 'pipe') or self.pipe is None:
             return {"error": "Model not loaded. Please check initialization logs."}
         
         # Parse the request payload
@@ -144,10 +157,20 @@ class EndpointHandler:
         except Exception as e:
             return {"error": f"Failed to parse request data: {str(e)}"}
         
+        # Extract parameters from the payload
+        parameters = {}
+        if "parameters" in payload and isinstance(payload["parameters"], dict):
+            # HF Inference Endpoints format: {"inputs": "prompt", "parameters": {...}}
+            parameters = payload["parameters"]
+        
         # Get the prompt from the payload
-        prompt_text = payload.get("prompt", "")
+        prompt_text = payload.get("inputs", "")
         if not prompt_text:
-            return {"error": "No prompt provided"}
+            # Try to get prompt from different fields for compatibility
+            prompt_text = payload.get("prompt", "")
+            
+        if not prompt_text:
+            return {"error": "No prompt provided. Please include 'inputs' or 'prompt' field."}
         
         # Apply child-content filtering to the prompt
         if contains_child_related_content(prompt_text):
@@ -155,17 +178,19 @@ class EndpointHandler:
         
         # 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", ""))
+        # Use negative_prompt from parameters or payload, fall back to config
+        negative_prompt = parameters.get("negative_prompt", payload.get("negative_prompt", self.cfg.get("negative_prompt", "")))
+        
+        # Get dimensions from config (default to 1024x768 if not specified)
+        width = int(self.cfg.get("width", 1024))
+        height = int(self.cfg.get("height", 768))
         
-        # 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)))
+        # Other generation parameters
+        inference_steps = int(parameters.get("inference_steps", payload.get("inference_steps", self.cfg.get("inference_steps", 30))))
+        guidance_scale = float(parameters.get("guidance_scale", 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)))
+        seed = int(parameters.get("seed", payload.get("seed", random.randint(0, MAX_SEED))))
         generator = torch.Generator(self.pipe.device).manual_seed(seed)
         
         try:
@@ -184,11 +209,56 @@ class EndpointHandler:
             # 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}
+            # Return the response formatted for Hugging Face Inference Endpoints
+            return [{"generated_image": 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} 
+            return {"error": error_message}
+
+# For local testing without HF Inference Endpoints
+if __name__ == "__main__":
+    import argparse
+    import uvicorn
+    from fastapi import FastAPI, Request
+    from fastapi.responses import JSONResponse
+    
+    # Parse command-line arguments
+    parser = argparse.ArgumentParser(description="Run the text-to-image API locally")
+    parser.add_argument("--port", type=int, default=8000, help="Port to run the server on")
+    parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to run the server on")
+    args = parser.parse_args()
+    
+    # Create FastAPI app
+    app = FastAPI(title="Text-to-Image API with Content Filtering")
+    
+    # Initialize the handler
+    handler = EndpointHandler()
+    
+    @app.get("/")
+    async def read_root():
+        """Health check endpoint."""
+        return {"status": "ok", "message": "Text-to-Image API is running"}
+    
+    @app.post("/")
+    async def generate_image(request: Request):
+        """Main inference endpoint."""
+        try:
+            body = await request.json()
+            result = handler(body)
+            
+            if "error" in result:
+                return JSONResponse(status_code=500, content={"error": result["error"]})
+                
+            return result
+        except Exception as e:
+            return JSONResponse(
+                status_code=500,
+                content={"error": f"Failed to process request: {str(e)}"}
+            )
+    
+    # Run the server
+    print(f"Starting server on http://{args.host}:{args.port}")
+    uvicorn.run(app, host=args.host, port=args.port) 

requirements.txt

@@ -1,9 +1,11 @@
-fastapi
-uvicorn[standard]
-diffusers
-transformers
-torch
-accelerate
-huggingface_hub
-pillow
-safetensors
+fastapi>=0.95.0
+uvicorn>=0.22.0
+torch>=2.0.0
+diffusers>=0.19.0
+transformers>=4.30.0
+accelerate>=0.20.0
+huggingface_hub>=0.16.0
+pydantic>=1.10.0
+Pillow>=9.0.0
+scipy>=1.10.0
+safetensors>=0.3.0

src/Procfile

@@ -1 +0,0 @@
-web: uvicorn app:app --host 0.0.0.0 --port $PORT

src/app.conf

@@ -1,11 +0,0 @@
-{
-    "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
-}

src/app.py

@@ -1,145 +0,0 @@
-import os
-import json
-import random
-import re
-import base64
-from io import BytesIO
-
-from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
-from PIL import Image
-
-import torch
-from huggingface_hub import snapshot_download
-from diffusers import (
-    AutoencoderKL,
-    StableDiffusionXLPipeline,
-    EulerAncestralDiscreteScheduler,
-    DPMSolverSDEScheduler
-)
-from diffusers.models.attention_processor import AttnProcessor2_0
-
-# Global constants
-MAX_SEED = 12211231
-NUM_IMAGES_PER_PROMPT = 1
-USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
-
-# Load configuration from app.conf
-with open("app.conf", "r") as f:
-    cfg = json.load(f)
-
-# --- 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."""
-    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."""
-    return bool(child_related_regex.search(prompt))
-
-# --- Model Pipeline Loading ---
-def load_pipeline_and_scheduler():
-    clip_skip = cfg.get("clip_skip", 0)
-    
-    # Download model files from Hugging Face Hub
-    ckpt_dir = snapshot_download(repo_id=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=cfg.get("use_safetensors", True),
-        variant="fp16"
-    )
-    pipe = pipe.to("cuda")
-    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(cfg.get("sampler", "DPM++ SDE Karras"))
-    
-    # Adjust the text encoder layers if needed using clip_skip
-    pipe.text_encoder.config.num_hidden_layers -= (clip_skip - 1)
-    
-    if USE_TORCH_COMPILE:
-        pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
-        print("Model Compiled!")
-    return pipe
-
-# Load the model pipeline once at startup
-pipe = load_pipeline_and_scheduler()
-
-# --- Utility Function: Convert PIL Image to Base64 ---
-def pil_image_to_base64(img: Image.Image) -> str:
-    buffered = BytesIO()
-    img.convert("RGB").save(buffered, format="WEBP", quality=90)
-    return base64.b64encode(buffered.getvalue()).decode("utf-8")
-
-# --- FastAPI Application Setup ---
-app = FastAPI(title="Text-to-Image API with Content Filtering")
-
-class GenerateRequest(BaseModel):
-    prompt: str
-
-@app.get("/")
-async def read_root():
-    return {"message": "Text-to-Image API with content filtering is running."}
-
-@app.post("/generate")
-async def generate(req: GenerateRequest):
-    # Apply child-content filtering to the prompt
-    prompt_text = req.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 = cfg.get("prompt", "{prompt}").replace("{prompt}", prompt_text)
-    # Use negative_prompt if provided; otherwise, default to empty string
-    negative_prompt = cfg.get("negative_prompt", "")
-    width = cfg.get("width", 1024)
-    height = cfg.get("height", 768)
-    inference_steps = cfg.get("inference_steps", 30)
-    guidance_scale = cfg.get("guidance_scale", 7)
-    
-    # Randomize the seed for generation
-    seed = random.randint(0, MAX_SEED)
-    generator = torch.Generator(pipe.device).manual_seed(seed)
-    
-    try:
-        outputs = 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])
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Image generation failed: {e}")
-    
-    return {"image_base64": img_base64, "seed": seed}
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run("app:app", host="0.0.0.0", port=8000)

src/requirements.txt

@@ -1,7 +0,0 @@
-fastapi
-uvicorn
-torch
-diffusers
-huggingface_hub
-pydantic
-Pillow