app.py

import gradio as gr
import numpy as np
import random 
# import spaces #[uncomment to use ZeroGPU]
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler, DPMSolverMultistepScheduler, UniPCMultistepScheduler, PNDMScheduler
from ultralytics import YOLO
from huggingface_hub import hf_hub_download
import torch 
import os
from PIL import Image, ImageFilter, ImageOps
from huggingface_hub import login, hf_hub_download
from gradio.themes import Default  # For theming


if "HF_TOKEN" in os.environ:
    login(os.environ["HF_TOKEN"])
else:
    raise ValueError("HF_TOKEN not found in environment variables. Please set it in Space settings.")

device = "cuda" if torch.cuda.is_available() else "cpu"

repo_id = "DreamingOracle/Quagmaform_alpha-1"  
filename = "DPS_Quagmaform_Alpha1.safetensors"  
model_path = hf_hub_download(repo_id=repo_id, filename=filename)
if torch.cuda.is_available(): 
    torch_dtype = torch.float16
else: 
    torch_dtype = torch.float32 

pipe = StableDiffusionPipeline.from_single_file(model_path, torch_dtype=torch_dtype) 
pipe = pipe.to(device) 

# Scheduler mapping (name to class)
SCHEDULERS = {
    "PNDM": PNDMScheduler,
    "Euler": EulerDiscreteScheduler,
    "DPM++ 2M Karras": DPMSolverMultistepScheduler,
    "UniPC": UniPCMultistepScheduler,
}

# Download ADetailer model if not present
adetailer_model_path = "face_yolov8n.pt"
if not os.path.exists(adetailer_model_path):
    hf_hub_download(repo_id="Bingsu/adetailer", filename="face_yolov8n.pt", local_dir=".")

adetailer_model = YOLO(adetailer_model_path)

MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 1024

# Purple theme
custom_theme = Default(primary_hue="purple")

def infer( 
    prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, scheduler_name="PNDM", save_format="png", progress=gr.Progress(track_tqdm=True),): 
    if randomize_seed: 
        seed = random.randint(0, MAX_SEED) 
    generator = torch.Generator().manual_seed(seed) 
    
    # Set scheduler dynamically
    scheduler_class = SCHEDULERS.get(scheduler_name, PNDMScheduler)
    pipe.scheduler = scheduler_class.from_config(pipe.scheduler.config)
    
    image = pipe( 
        prompt=prompt, negative_prompt=negative_prompt, guidance_scale=guidance_scale, num_inference_steps=num_inference_steps, width=width, height=height, generator=generator, 
    ).images[0] 
    
    # ---------------------------
    # ADetailer post-processing for face enhancement (with padding + soft blend)
    # ---------------------------
    try:
        results = adetailer_model(image)
        if results and len(results) and getattr(results[0], "boxes", None):
            for box in results[0].boxes:
                x1, y1, x2, y2 = map(int, box.xyxy[0])

                w = max(1, x2 - x1)
                h = max(1, y2 - y1)
                pad = int(max(10, 0.18 * max(w, h)))
                x1p = max(0, x1 - pad)
                y1p = max(0, y1 - pad)
                x2p = min(image.width, x2 + pad)
                y2p = min(image.height, y2 + pad)

                face = image.crop((x1p, y1p, x2p, y2p))

                fw, fh = face.size
                fw8 = max(8, (fw // 8) * 8)
                fh8 = max(8, (fh // 8) * 8)
                if (fw8, fh8) != (fw, fh):
                    face = face.resize((fw8, fh8), Image.LANCZOS)

                mask = Image.new("L", face.size, 255)

                blur_radius = max(4, int(min(face.size) / 10))
                paste_mask = mask.filter(ImageFilter.GaussianBlur(radius=blur_radius))

                inpaint_result = pipe(
                    prompt=prompt + ", high detail face",
                    image=face,
                    mask_image=mask,
                    strength=0.45,
                    num_inference_steps=20,
                    guidance_scale=7.5,
                    generator=generator
                ).images[0]

                if paste_mask.mode != "L":
                    paste_mask = paste_mask.convert("L")

                image.paste(inpaint_result, (x1p, y1p), paste_mask)
    except Exception as e:
        print("ADetailer post-process failed:", e)

    output_path = f"generated_image.{save_format}"
    image.save(output_path, format=save_format.upper())
    return image, seed 

examples = [ 
    "photorealistic portrait of a young woman, cinematic rim lighting, soft golden hour backlight, detailed skin pores, realistic eyelashes, 85mm lens, shallow depth of field, ultra-detailed, high dynamic range, film grain, detailed, 8k", 
    "head helmet portrait of a futuristic armored soldier, worn brushed metal armor with neon blue accents, realistic cloth under-armor, weathering and scratches, volumetric rim light, cinematic pose, high detail, photoreal", 
    "arctic mountain in snow, insulated modules, panorama view, blowing snow, cold blue light, realistic snow accumulation, high detail",] 

# Updated CSS 12826
css = """
#col-container { margin: 0 auto; max-width: 640px;}
#community-row {justify-content: center; gap: 30px;}
"""

with gr.Blocks(css=css, theme=custom_theme) as demo: 
    with gr.Column(elem_id="col-container"): 
        gr.Markdown("# DPS-Quagmaform AI txt2img") 
        with gr.Row(): 
            prompt = gr.Text( 
                label="Prompt", show_label=False, max_lines=1, placeholder="Enter your prompt", container=False, 
            ) 
            run_button = gr.Button("Run", scale=0, variant="primary") 
        result = gr.Image(label="Result", show_label=False) 
        with gr.Accordion("Advanced Settings", open=False): 
            negative_prompt = gr.Text( 
                label="Negative prompt", max_lines=1, placeholder="Enter a negative prompt", visible=True, 
            ) 
            seed = gr.Slider( 
                label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0, 
            ) 
            randomize_seed = gr.Checkbox(label="Randomize seed", value=True) 
            with gr.Row(): 
                width = gr.Slider( 
                    label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=768,
                ) 
                height = gr.Slider( 
                    label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024, 
                ) 
            with gr.Row(): 
                guidance_scale = gr.Slider( 
                    label="Guidance scale", minimum=0.0, maximum=10.0, step=0.1, value=5, 
                ) 
                num_inference_steps = gr.Slider( 
                    label="Number of inference steps", minimum=1, maximum=50, step=1, value=22, 
                ) 
            scheduler = gr.Dropdown( 
                label="Sampler/Scheduler", 
                choices=list(SCHEDULERS.keys()), 
                value="PNDM", 
                info="Change this setting for better quality in some situations" 
            ) 
            save_format = gr.Dropdown( 
                choices=["png", "jpg"], value="png", label="Select Output Format" 
            ) 
        gr.Examples(examples=examples, inputs=[prompt]) 
        
        # Community
        gr.Markdown("### Community")
        with gr.Row(elem_id="community-row"):
            
            gr.Button("Join Discord 💬", link="https://discord.gg/deepspace", variant="primary")
            gr.Button("Telegram En Español 📱", link="https://t.me/DeepSpaceHispano", variant="primary")
        
        gr.on( 
            triggers=[run_button.click, prompt.submit], 
            fn=infer, 
            inputs=[ 
                prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, scheduler, save_format, 
            ], 
            outputs=[result, seed], 
        ) 
if __name__ == "__main__": 
    demo.launch()