Spaces:
Runtime error
Runtime error
| """Inference backend for the DiffusionDB SD 1.5 LoRA. | |
| A Gradio app whose only job is to expose generation as an API. The UI Space | |
| (spaces/whosouravsharma/diffusiondb-sd15-lora) calls it with gradio_client; | |
| the minimal interface here is also usable directly for debugging. | |
| This is the Space that needs the GPU. Weights are pulled from the model repo | |
| at first request, not baked into the image. | |
| """ | |
| import os | |
| import gradio as gr | |
| from src.pipeline import LORA_CHECKPOINT, LORA_REPO, ModelRunner, pick_device | |
| from src.tracing import ( | |
| TRACING_ACTIVE, | |
| configure_logging, | |
| continue_trace, | |
| logger, | |
| ) | |
| configure_logging() | |
| # Loaded on first request rather than at import, so the Space becomes | |
| # reachable immediately and a slow weight download cannot kill startup. | |
| _runner = None | |
| def get_runner() -> ModelRunner: | |
| global _runner | |
| if _runner is None: | |
| _runner = ModelRunner() | |
| return _runner | |
| def model_is_warm() -> bool: | |
| return _runner is not None | |
| # Status strings for the visible UI. The cold path is genuinely slow -- the | |
| # container has to wake, then torch/diffusers import, then ~9 s of weight | |
| # loading before a single step runs -- so say so rather than showing a | |
| # spinner that looks indistinguishable from a hang. | |
| STATUS_IDLE = "Ready." | |
| STATUS_COLD = ( | |
| "⏳ **Starting server…** waking the GPU and loading the model. " | |
| "The first image takes around 40 seconds; later ones take a few." | |
| ) | |
| STATUS_WARM = "⏳ **Generating…**" | |
| def status_on_click() -> str: | |
| """Fires the instant the button is pressed, before the GPU job queues.""" | |
| return STATUS_WARM if model_is_warm() else STATUS_COLD | |
| def status_on_load() -> str: | |
| """Fires when someone opens the page.""" | |
| if model_is_warm(): | |
| return STATUS_IDLE | |
| return ( | |
| "Model not loaded yet — the first generation also starts the server " | |
| "and takes around 40 seconds." | |
| ) | |
| def generate(prompt, negative_prompt, steps, guidance, lora_scale, seed, | |
| request: gr.Request = None): | |
| """One image per call. Returns (image, seed_used). | |
| Deliberately single-image: gradio_client handles a plain Image return | |
| far more predictably than a Gallery, and the UI can call again for more. | |
| `request` is supplied by Gradio from the type annotation and is not a UI | |
| input, so the published /generate signature stays at six arguments. It | |
| exists only to pick up the caller's trace headers. | |
| """ | |
| try: | |
| # When the UI Space sent trace headers this run becomes a child of | |
| # its run, so one LangSmith trace covers UI and backend -- and the | |
| # visitor's location, carried in baggage, lands on this run too. | |
| with continue_trace(request): | |
| result = get_runner().generate( | |
| prompt=prompt, | |
| negative_prompt=negative_prompt, | |
| steps=int(steps), | |
| guidance=float(guidance), | |
| lora_scale=float(lora_scale), | |
| seed=None if int(seed) < 0 else int(seed), | |
| count=1, | |
| ) | |
| except ValueError as error: | |
| raise gr.Error(str(error)) | |
| except Exception as error: | |
| # The gr.Error string is all the caller sees; without this the | |
| # traceback is lost and the Space logs show nothing useful. | |
| logger.exception("request failed") | |
| raise gr.Error(f"Generation failed: {error}") | |
| return result.images[0], result.seed | |
| with gr.Blocks(title="DiffusionDB SD 1.5 LoRA — inference API") as demo: | |
| gr.Markdown( | |
| f"## Inference API\n" | |
| f"Backend for [`diffusiondb-sd15-lora`]" | |
| f"(https://huggingface.co/spaces/whosouravsharma/diffusiondb-sd15-lora).\n\n" | |
| f"Model `{LORA_REPO}` @ `{LORA_CHECKPOINT}` · device `{pick_device()}` · " | |
| f"tracing `{'on' if TRACING_ACTIVE else 'off'}`\n\n" | |
| f"Call it with `gradio_client`:\n" | |
| f"```python\n" | |
| f"from gradio_client import Client\n" | |
| f"c = Client('whosouravsharma/diffusiondb-sd15-lora-inference')\n" | |
| f"image, seed = c.predict('a steampunk owl', '', 25, 7.5, 1.0, 42,\n" | |
| f" api_name='/generate')\n" | |
| f"```" | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| prompt = gr.Textbox(label="prompt", lines=2) | |
| negative = gr.Textbox(label="negative_prompt", lines=1) | |
| steps = gr.Number(value=25, label="steps", precision=0) | |
| guidance = gr.Number(value=7.5, label="guidance") | |
| lora_scale = gr.Number(value=1.0, label="lora_scale") | |
| seed = gr.Number(value=-1, label="seed", precision=0) | |
| run = gr.Button("Generate", variant="primary") | |
| with gr.Column(): | |
| # Initial value assumes cold, because at import time it always | |
| # is; demo.load corrects it to "Ready." on a warm server. The | |
| # reverse default would flash a false "Ready." at every cold | |
| # visitor before the load event lands. | |
| status = gr.Markdown(status_on_load()) | |
| image = gr.Image(label="image", type="pil") | |
| used_seed = gr.Number(label="seed_used", interactive=False) | |
| # Two independent listeners on the same button. The cosmetic one uses | |
| # queue=False so it paints immediately instead of waiting behind the GPU | |
| # job it is announcing. The API contract of /generate is untouched: same | |
| # six inputs, same two outputs, so gradio_client callers are unaffected. | |
| # api_name=False keeps these cosmetic events out of the published API, | |
| # so gradio_client sees /generate and nothing else. | |
| run.click(status_on_click, outputs=status, queue=False, api_name=False) | |
| generate_event = run.click( | |
| generate, | |
| inputs=[prompt, negative, steps, guidance, lora_scale, seed], | |
| outputs=[image, used_seed], | |
| api_name="generate", | |
| ) | |
| generate_event.then( | |
| lambda: STATUS_IDLE, outputs=status, queue=False, api_name=False | |
| ) | |
| demo.load(status_on_load, outputs=status, api_name=False) | |
| if __name__ == "__main__": | |
| demo.queue(max_size=12).launch( | |
| server_name="0.0.0.0", | |
| server_port=int(os.environ.get("PORT", 7860)), | |
| ) | |