Hugging Face's logo

Spaces:
rickveloper
/
Vibeforge
Runtime error

App Files Community
Vibeforge
File size: 4,921 Bytes 
ec03850
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cde0214
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
 
import os, random, re, torch
from typing import List, Tuple
from PIL import Image, ImageDraw, ImageFont
import gradio as gr
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler

# --------------------
# Config
# --------------------
MODEL_ID = os.getenv("MODEL_ID", "runwayml/stable-diffusion-v1-5")
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32

# Simple prompt guardrail (blocks obvious NSFW attempts)
NSFW_TERMS = [
    r"\bnsfw\b", r"\bnude\b", r"\bnudity\b", r"\bsex\b", r"\bexplicit\b", r"\bporn\b",
    r"\bboobs\b", r"\bbutt\b", r"\bass\b", r"\bnsfw\b", r"\bnaked\b", r"\btits\b",
    r"\b18\+\b", r"\berotic\b", r"\bfetish\b"
]
NSFW_REGEX = re.compile("|".join(NSFW_TERMS), flags=re.IGNORECASE)

# --------------------
# Load pipeline
# --------------------
pipe = StableDiffusionPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=DTYPE
)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

if torch.cuda.is_available():
    pipe = pipe.to("cuda")
    pipe.enable_attention_slicing()
    pipe.enable_vae_slicing()
else:
    pipe = pipe.to("cpu")

# --------------------
# Helpers
# --------------------
def blocked_tile(reason: str, width=512, height=512) -> Image.Image:
    img = Image.new("RGB", (width, height), (20, 20, 24))
    draw = ImageDraw.Draw(img)
    text = f"BLOCKED\n{reason}"
    try:
        font = ImageFont.truetype("DejaVuSans-Bold.ttf", 28)
    except:
        font = ImageFont.load_default()
    tw, th = draw.multiline_textbbox((0,0), text, font=font)[2:]
    draw.multiline_text(((width - tw)//2, (height - th)//2), text, fill=(255,255,255), font=font, align="center")
    return img

def is_prompt_nsfw(prompt: str) -> bool:
    return bool(NSFW_REGEX.search(prompt or ""))

def generate(
    prompt: str,
    negative_prompt: str,
    steps: int,
    guidance: float,
    width: int,
    height: int,
    seed: int,
    batch_size: int
) -> Tuple[List[Image.Image], str]:
    if not prompt.strip():
        return [], "Add a prompt to get rolling."

    # Hard block obvious NSFW prompts before hitting the model
    if is_prompt_nsfw(prompt) or is_prompt_nsfw(negative_prompt or ""):
        img = blocked_tile("NSFW prompt detected")
        return [img], "Blocked: NSFW prompt."

    # Seed
    if seed < 0:
        seed = random.randint(0, 2**31 - 1)
    generator = torch.Generator(device=DEVICE).manual_seed(seed)

    out = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt or None,
        num_inference_steps=steps,
        guidance_scale=guidance,
        width=width,
        height=height,
        num_images_per_prompt=batch_size,
        generator=generator
    )

    images = out.images
    flags = getattr(out, "nsfw_content_detected", None)

    # If the underlying safety checker flags NSFW, block it (no blur)
    if flags:
        for i, flagged in enumerate(flags):
            if flagged:
                images[i] = blocked_tile("NSFW content flagged")

    msg = f"Seed: {seed} • Images: {len(images)}"
    if flags is not None:
        msg += f" • Flagged: {sum(1 for f in flags if f)}"
    return images, msg

# --------------------
# UI
# --------------------
with gr.Blocks(title="VibeForge — Clean Image Generator") as demo:
    gr.Markdown(
        """
# VibeForge ⚒️  
**Clean, creative image generation.**  
NSFW inputs are blocked. Keep it classy and go wild on style, lighting, composition, mood.
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="a cinematic photo of a vintage motorcycle by the ocean at sunset, golden hour, soft rim light, 50mm"
            )
            negative = gr.Textbox(label="Negative Prompt", placeholder="low quality, blurry, watermark, jpeg artifacts")
            with gr.Row():
                steps = gr.Slider(10, 50, value=28, step=1, label="Steps")
                guidance = gr.Slider(1.0, 12.0, value=7.5, step=0.1, label="CFG")
            with gr.Row():
                width = gr.Dropdown(choices=[384, 448, 512, 640, 768], value=512, label="Width")
                height = gr.Dropdown(choices=[384, 448, 512, 640, 768], value=512, label="Height")
            with gr.Row():
                seed = gr.Number(value=-1, label="Seed (-1 = random)", precision=0)
                batch = gr.Slider(1, 4, value=1, step=1, label="Batch")

            go = gr.Button("Generate", variant="primary")

        with gr.Column(scale=5):
            gallery = gr.Gallery(label="Output", columns=2, height=512)
            info = gr.Markdown()

    go.click(
        fn=generate,
        inputs=[prompt, negative, steps, guidance, width, height, seed, batch],
        outputs=[gallery, info]
    )

if __name__ == "__main__":
    demo.launch()