Blame for app.py

"""Personalized text-to-image website. Requires Sumith's hosted model weights."""
import os
import secrets
import threading
from pathlib import Path

import spaces
import gradio as gr

from prompts import make_prompt, STYLES

_pipeline = None
_lock = threading.Lock()


def load_pipeline():
    global _pipeline
    if _pipeline is not None:
        return _pipeline
    import torch
    from diffusers import StableDiffusionPipeline, DDIMScheduler
    from huggingface_hub import hf_hub_download

    local_path = os.getenv("MODEL_PATH", "")
    repo = os.getenv("MODEL_REPO_ID", "Sumith0102/PersonalisedAiAvatarGenerator")
    if local_path:
        path = Path(local_path)
    elif repo:
        path = Path(hf_hub_download(
            repo_id=repo,
            filename=os.getenv("MODEL_FILENAME", "model.ckpt"),
            token=os.getenv("HF_TOKEN") or None,
        ))
    else:
        raise gr.Error("The owner is connecting the personalized model. Please try again later.")
    if path.suffix == ".ckpt":
        from convert_checkpoint import convert
        import hashlib
        import tempfile
        cache_key = hashlib.sha256(str(path.resolve()).encode()).hexdigest()[:20]
        safe_path = Path(tempfile.gettempdir()) / f"sumith-{cache_key}.safetensors"
        if not safe_path.exists():
            convert(path, safe_path)
        path = safe_path
    elif path.suffix != ".safetensors":
        raise gr.Error("The configured model format is not supported.")
    pipeline = StableDiffusionPipeline.from_single_file(
        str(path),
        config=os.getenv("BASE_MODEL_ID", "stable-diffusion-v1-5/stable-diffusion-v1-5"),
        torch_dtype=torch.float16,
        token=os.getenv("HF_TOKEN") or None,
        load_safety_checker=True,
    )
    pipeline.scheduler = DDIMScheduler.from_config(pipeline.scheduler.config)
    pipeline.enable_attention_slicing()
    pipeline.enable_vae_slicing()
    pipeline.to("cuda")
    _pipeline = pipeline
    return pipeline


def prepare_prompt(scene, style, enhance=True):
    try:
        return make_prompt(scene, style, load_pipeline().tokenizer, enhance)
    except ValueError as exc:
        raise gr.Error(str(exc)) from None


@spaces.GPU(duration=60)
def generate(scene, style, seed, steps, enhance=True, progress=gr.Progress()):
    import torch
    prompt = prepare_prompt(scene, style, enhance)
    if seed is None:
        raise gr.Error("Enter a seed, or use -1 for a random image.")
    if isinstance(seed, bool) or int(seed) != seed or not -1 <= seed <= 2147483647:
        raise gr.Error("Choose a whole-number seed from 0 to 2147483647, or -1 for random.")
    seed = secrets.randbelow(2147483648) if seed == -1 else int(seed)
    steps = int(steps)
    if not 15 <= steps <= 35:
        raise gr.Error("Use between 15 and 35 generation steps.")
    with _lock:
        try:
            progress(0.05, desc="Loading the personalized model")
            pipe = load_pipeline()
            device = "cuda"
            rng = torch.Generator(device=device).manual_seed(seed)
            def step_callback(pipeline, index, timestep, kwargs):
                progress(0.1 + 0.85 * (index + 1) / steps, desc="Creating your portrait")
                return kwargs
            with torch.inference_mode():
                result = pipe(
                    prompt=prompt,
                    negative_prompt="blurry, low resolution, distorted face, duplicate person, watermark, text",
                    height=512, width=512, guidance_scale=7.5,
                    num_inference_steps=steps, generator=rng,
                    callback_on_step_end=step_callback,
                )
            if result.nsfw_content_detected and any(result.nsfw_content_detected):
                raise gr.Error("This image could not be displayed. Please try another scene.")
            return result.images[0], f"Seed: {seed} · 512 × 512 · AI-generated image"
        except gr.Error:
            raise
        except Exception as exc:
            # Do not return upstream paths, tokens, prompts, or credentials to visitors.
            print("Generation failed:", type(exc).__name__, flush=True)
            raise gr.Error("Generation is unavailable right now. Please try again later.") from None


CSS = """
.gradio-container {max-width: 1240px !important; margin: auto; background: #f5f7fc;}
#studio-header {padding: 28px 0 14px; border-bottom: 1px solid #cdd5e6; margin-bottom: 22px;}
#studio-header h1 {font-size: clamp(30px, 4vw, 48px); letter-spacing: -1.5px; color: #182448;}
#studio-header p {font-size: 16px; color: #475777;}
#generate {background: #304ffe; color: white; border: 0; min-height: 52px; font-size: 17px;}
#portrait {min-height: 460px; background: #e9eef9; border-radius: 18px;}
#footer {font-size: 14px; color: #566581; margin-top: 20px;}
"""

with gr.Blocks(title="Sumith Portrait Studio",
               analytics_enabled=False, delete_cache=(3600, 3600)) as demo:
    gr.Markdown("# Sumith Portrait Studio\nDescribe a scene. Create an AI portrait of Sumith.", elem_id="studio-header")
    with gr.Row():
        with gr.Column(scale=5):
            scene = gr.Textbox(label="Your scene", lines=5, max_lines=8,
                placeholder="Standing beside a lake in the Canadian mountains at sunrise")
            style = gr.Radio(list(STYLES), value="Natural", label="Look")
            enhance = gr.Checkbox(value=True, label="Add visual detail to my prompt")
            gr.Markdown("A lightweight prompt assistant adds lighting and composition cues where there is room. "
                        "Your scene stays intact. No extra AI service is used.")
            preview = gr.Button("Preview generation prompt", variant="secondary")
            final_prompt = gr.Textbox(label="Exact generation prompt", lines=4, interactive=False,
                                      placeholder="Preview your prompt, or create a portrait to see it here.")
            gr.Markdown("The model has a short text limit. Optional details only fill the remaining space; "
                        "overlong scenes are rejected, never silently cut off. Turn enhancement off for just your scene and the identity cue.")
            with gr.Accordion("Fine-tune the result", open=False):
                seed = gr.Number(value=-1, precision=0, label="Seed · -1 for a new variation")
                steps = gr.Slider(15, 35, value=24, step=1, label="Generation steps")
            submit = gr.Button("Create portrait", variant="primary", elem_id="generate")
            gr.Examples([
                ["Walking through Old Montréal on a sunny afternoon", "Natural"],
                ["In a tailored navy suit, soft studio lighting, professional headshot", "Studio"],
                ["An explorer beside a mountain lake at sunrise", "Cinematic"],
            ], inputs=[scene, style], label="Try a scene", cache_examples=False)
        with gr.Column(scale=6):
            portrait = gr.Image(label="Your portrait", type="pil", interactive=False,
                                buttons=["download"], elem_id="portrait")
            details = gr.Markdown("Your generated portrait will appear here.")
    gr.Markdown(
        "Built by Sumith Reddy · DreamBooth / Stable Diffusion 1.5  \n"
        "Images are AI-generated and may vary in likeness. Prompts and results are processed by the hosting service. "
        "The app keeps temporary output files for up to about two hours.  \n"
        "[Back to portfolio](https://thanu-sumith.github.io/sumith-portfolio/)", elem_id="footer")
    preview.click(prepare_prompt, inputs=[scene, style, enhance], outputs=final_prompt,
                  api_name="preview_prompt")
    # Clear stale previews when the request changes.
    for control in [scene, style, enhance]:
        control.change(lambda: "", outputs=final_prompt, queue=False, api_name=False)
    submit.click(prepare_prompt, inputs=[scene, style, enhance], outputs=final_prompt,
                 api_name=False).success(
        generate, inputs=[scene, style, seed, steps, enhance], outputs=[portrait, details],
        concurrency_limit=1, api_name="generate_portrait")

# Prepare the model before ZeroGPU forks a worker. Its CUDA emulation supports
# model placement at startup; real GPU computation happens only in generate().
if __name__ == "__main__":
    load_pipeline()
    demo.queue(max_size=4, default_concurrency_limit=1).launch(
        css=CSS, theme=gr.themes.Soft(primary_hue="indigo", neutral_hue="slate"),
        server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")),
        show_error=False, max_file_size="1mb",
    )