chore: init commit

.python-version

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

src/Procfile

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

src/app.conf

@@ -0,0 +1,13 @@
+{
+"cover_path_": ["hentai-waianiv6-card0.webp", "hentai-waianiv6-card1.webp", "hentai-waianiv6-card2.webp", "hentai-waianiv6-card3.webp", "hentai-waianiv6-card4.webp", "hentai-waianiv6-card5.webp", "hentai-waianiv6-card6.webp", "hentai-waianiv6-card7.webp"],
+"cover_path": ["hentai-waianiv6-card5.webp"],
+"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

@@ -0,0 +1,145 @@
+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

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