"""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", )