app.py

import os

os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

import spaces
import gradio as gr
import numpy as np
import random
import time
import torch
from diffusers.pipelines.glm_image import GlmImagePipeline
from PIL import Image

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

pipe = GlmImagePipeline.from_pretrained(
    "zai-org/GLM-Image",
    torch_dtype=torch.bfloat16,
).to("cuda")


def _gpu_duration(
    prompt,
    input_images=None,
    seed=42,
    randomize_seed=False,
    width=1024,
    height=1024,
    num_inference_steps=50,
    guidance_scale=1.5,
    *args,
    **kwargs,
):
    steps = int(num_inference_steps or 50)
    return min(240, 90 + steps * 3)


@spaces.GPU(duration=_gpu_duration)
def infer(
    prompt,
    input_images=None,
    seed=42,
    randomize_seed=False,
    width=1024,
    height=1024,
    num_inference_steps=50,
    guidance_scale=1.5,
    progress=gr.Progress(track_tqdm=True),
):
    if not prompt or not str(prompt).strip():
        raise gr.Error("Please enter a prompt.")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    width = max(256, min(MAX_IMAGE_SIZE, (int(width) // 32) * 32))
    height = max(256, min(MAX_IMAGE_SIZE, (int(height) // 32) * 32))

    generator = torch.Generator(device="cuda").manual_seed(int(seed))

    image_list = None
    if input_images is not None and len(input_images) > 0:
        image_list = []
        for item in input_images:
            img = item[0] if isinstance(item, tuple) else item
            if isinstance(img, str):
                img = Image.open(img).convert("RGB")
            elif isinstance(img, Image.Image):
                img = img.convert("RGB")
            else:
                continue
            image_list.append(img)
        if not image_list:
            image_list = None

    pipe_kwargs = {
        "prompt": prompt,
        "height": height,
        "width": width,
        "num_inference_steps": int(num_inference_steps),
        "guidance_scale": float(guidance_scale),
        "generator": generator,
    }
    if image_list is not None:
        pipe_kwargs["image"] = image_list

    started = time.perf_counter()
    image = pipe(**pipe_kwargs).images[0]
    elapsed = time.perf_counter() - started
    print(
        f"infer_seconds={elapsed:.1f} size={width}x{height} "
        f"steps={num_inference_steps} i2i={image_list is not None}"
    )
    return image, seed


def update_dimensions_from_image(image_list):
    if image_list is None or len(image_list) == 0:
        return 1024, 1024

    item = image_list[0]
    img = item[0] if isinstance(item, tuple) else item
    if isinstance(img, str):
        img = Image.open(img)

    img_width, img_height = img.size
    aspect_ratio = img_width / float(img_height)

    if aspect_ratio >= 1:
        new_width = 1024
        new_height = int(1024 / aspect_ratio)
    else:
        new_height = 1024
        new_width = int(1024 * aspect_ratio)

    new_width = max(256, min(MAX_IMAGE_SIZE, round(new_width / 32) * 32))
    new_height = max(256, min(MAX_IMAGE_SIZE, round(new_height / 32) * 32))
    return new_width, new_height


css = """
#col-container {
    margin: 0 auto;
    max-width: 1200px;
}
.gallery-container img {
    object-fit: contain;
}
"""

examples = [
    "A vintage coffee shop chalkboard menu, dark green board, white chalk lettering, text: 'Morning Blend $4', 'Cinnamon Roll $3', 'Iced Latte $5', steam and coffee beans in the corners",
    "A clean modern recipe card on cream paper, bold title 'Raspberry Mousse Cake', ingredient list 'Flour 150g', 'Eggs 3', 'Sugar 120g', 'Raspberry puree 200g', step boxes with short instructions, mint leaves and raspberries",
    "A concert poster for 'NIGHT OWL JAZZ', Friday October 10, The Blue Note Club, midnight set, gold foil serif type on deep navy, saxophone silhouette, city lights",
]

with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """# GLM-Image Pro
Hybrid autoregressive + DiT image generator from [zai-org/GLM-Image](https://huggingface.co/zai-org/GLM-Image), with Glyph Encoder for in-image text.

- **Text-to-image**: leave the gallery empty
- **Image-to-image**: upload one or more images, then describe the edit
- Put any text that should appear in the image in quotation marks
"""
        )

        with gr.Row():
            with gr.Column():
                prompt = gr.Text(
                    label="Prompt",
                    show_label=False,
                    max_lines=6,
                    placeholder="Describe the image. Put rendered text in quotes, e.g. title 'Harbor Market'.",
                    container=False,
                )
                run_button = gr.Button("Generate", variant="primary")

                with gr.Accordion("Input image(s) for image-to-image", open=False):
                    input_images = gr.Gallery(
                        label="Input Image(s)",
                        type="pil",
                        columns=3,
                        rows=1,
                        elem_classes="gallery-container",
                    )

                with gr.Accordion("Advanced settings", open=False):
                    seed = gr.Slider(
                        label="Seed",
                        minimum=0,
                        maximum=MAX_SEED,
                        step=1,
                        value=42,
                    )
                    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=1024,
                            info="Must be a multiple of 32",
                        )
                        height = gr.Slider(
                            label="Height",
                            minimum=256,
                            maximum=MAX_IMAGE_SIZE,
                            step=32,
                            value=1024,
                            info="Must be a multiple of 32",
                        )
                    with gr.Row():
                        num_inference_steps = gr.Slider(
                            label="Inference steps",
                            minimum=1,
                            maximum=100,
                            step=1,
                            value=30,
                        )
                        guidance_scale = gr.Slider(
                            label="Guidance scale",
                            minimum=0.0,
                            maximum=10.0,
                            step=0.1,
                            value=1.5,
                        )

                gr.Examples(examples=examples, inputs=[prompt])

            with gr.Column():
                result = gr.Image(label="Result", show_label=False)

    input_images.upload(
        fn=update_dimensions_from_image,
        inputs=[input_images],
        outputs=[width, height],
    )

    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[
            prompt,
            input_images,
            seed,
            randomize_seed,
            width,
            height,
            num_inference_steps,
            guidance_scale,
        ],
        outputs=[result, seed],
    )

if __name__ == "__main__":
    demo.launch()