Spaces:
Sleeping
Sleeping
| """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 | |
| 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", | |
| ) | |