Z-Image

Running on Zero

Update app.py

b5f9b50 verified 3 months ago

14.6 kB

import os
import torch
import gradio as gr
import spaces
import diffusers

from diffusers import ZImagePipeline


# ============================================================
# VERSION CHECK
# ============================================================

print("=" * 60)
print("VERSION CHECK")
print("=" * 60)
print("Torch:", torch.__version__)
print("Diffusers:", diffusers.__version__)
print(
    "Transformers:",
    __import__("transformers").__version__,
)
print(
    "HuggingFace Hub:",
    __import__("huggingface_hub").__version__,
)
print(
    "Gradio:",
    gr.__version__,
)
print(
    "CUDA available:",
    torch.cuda.is_available(),
)

print(
    "Has ZImagePipeline:",
    hasattr(diffusers, "ZImagePipeline"),
)

print("=" * 60)


# ============================================================
# CHECK Z-IMAGE PIPELINE
# ============================================================

if not hasattr(diffusers, "ZImagePipeline"):
    raise RuntimeError(
        "ZImagePipeline is not available. "
        "Please install diffusers==0.36.0 or newer."
    )


# ============================================================
# LOAD Z-IMAGE
# ============================================================

print("Loading Z-Image-Turbo pipeline...")

pipe = ZImagePipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=False,
)

pipe.to("cuda")

print("Z-Image pipeline loaded!")


# ============================================================
# WATERMARK MODULE
# ============================================================

from watermark import (
    preview,
    remove_image,
    process_video,
)

from watermark.video import get_video_info


# ============================================================
# Z-IMAGE GENERATION
# ============================================================

@spaces.GPU(duration=60)
def generate_image(
    prompt: str,
    height: int,
    width: int,
    steps: int,
    seed: int,
    randomize_seed: bool,
    progress=gr.Progress(track_tqdm=True),
):
    """
    Generate an image with Z-Image-Turbo.
    """

    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    if randomize_seed:
        seed = torch.randint(
            0,
            2**32 - 1,
            (1,),
        ).item()

    seed = int(seed)

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

    print(
        f"Generating image: "
        f"{width}x{height}, "
        f"steps={steps}, "
        f"seed={seed}"
    )

    image = pipe(
        prompt=prompt,
        height=int(height),
        width=int(width),
        num_inference_steps=int(steps),
        guidance_scale=0.0,
        generator=generator,
    ).images[0]

    return image, seed


# ============================================================
# WATERMARK PREVIEW
# ============================================================

@spaces.GPU(duration=30)
def wm_preview(
    image,
    detection_prompt: str = "watermark",
    max_bbox_percent: float = 30,
):
    """
    Detect likely watermark regions.
    """

    if image is None:
        raise gr.Error("Please upload an image.")

    return preview(
        image,
        float(max_bbox_percent),
        detection_prompt or "watermark",
    )


# ============================================================
# WATERMARK IMAGE REMOVAL
# ============================================================

@spaces.GPU(duration=45)
def wm_remove_image(
    image,
    detection_prompt: str = "watermark",
    max_bbox_percent: float = 30,
    transparent: bool = False,
):
    """
    Detect and remove watermark regions from an image.
    """

    if image is None:
        raise gr.Error("Please upload an image.")

    return remove_image(
        image,
        float(max_bbox_percent),
        detection_prompt or "watermark",
        bool(transparent),
    )


# ============================================================
# WATERMARK VIDEO REMOVAL
# ============================================================

@spaces.GPU(duration=60)
def wm_remove_video(
    video_path: str,
    detection_prompt: str = "watermark",
    max_bbox_percent: float = 30,
    detection_skip: int = 5,
    fade_in: int = 2,
    fade_out: int = 2,
    progress=gr.Progress(),
):
    """
    Detect and remove watermarks from a video.
    """

    if not video_path:
        raise gr.Error(
            "Please upload a video."
        )

    # --------------------------------------------------------
    # Inspect video
    # --------------------------------------------------------

    try:

        duration, _, _, _, _ = get_video_info(
            video_path
        )

        max_seconds = float(
            os.getenv(
                "WATERMARK_MAX_VIDEO_SECONDS",
                "30",
            )
        )

        if duration > max_seconds + 0.25:

            raise gr.Error(
                f"Video is {duration:.1f}s. "
                f"Maximum is "
                f"{max_seconds:.0f}s "
                f"for this ZeroGPU build. "
                f"Please trim the video."
            )

    except gr.Error:
        raise

    except Exception as e:

        raise gr.Error(
            f"Unable to inspect video: {e}"
        )

    # --------------------------------------------------------
    # Process video
    # --------------------------------------------------------

    return process_video(
        video_path,
        detection_prompt or "watermark",
        float(max_bbox_percent),
        max(1, int(detection_skip)),
        max(0, int(fade_in)),
        max(0, int(fade_out)),
        progress=progress,
    )


# ============================================================
# GRADIO UI
# ============================================================

with gr.Blocks(
    title="Z-Image + Watermark Remover V5.2"
) as demo:

    gr.Markdown(
        """
# 🎨 Z-Image + Watermark Remover V5.2
Z-Image-Turbo + Florence-2 + LaMa
支持：
- 🎨 AI 文生图
- 🧹 图片去水印
- 🎬 短视频去水印
- 🔌 Gradio API
"""
    )

    # ========================================================
    # TABS
    # ========================================================

    with gr.Tabs():

        # ====================================================
        # Z-IMAGE
        # ====================================================

        with gr.Tab("🎨 Z-Image"):

            with gr.Row():

                with gr.Column():

                    prompt = gr.Textbox(
                        label="Prompt",
                        placeholder=(
                            "Describe the image..."
                        ),
                        lines=5,
                    )

                    with gr.Row():

                        height = gr.Slider(
                            minimum=512,
                            maximum=2048,
                            value=1024,
                            step=64,
                            label="Height",
                        )

                        width = gr.Slider(
                            minimum=512,
                            maximum=2048,
                            value=1024,
                            step=64,
                            label="Width",
                        )

                    with gr.Row():

                        steps = gr.Slider(
                            minimum=1,
                            maximum=30,
                            value=9,
                            step=1,
                            label="Inference Steps",
                        )

                        seed = gr.Number(
                            value=42,
                            precision=0,
                            label="Seed",
                        )

                    randomize = gr.Checkbox(
                        value=True,
                        label="Randomize seed",
                    )

                    generate_btn = gr.Button(
                        "Generate",
                        variant="primary",
                    )

                with gr.Column():

                    output_image = gr.Image(
                        label="Generated Image",
                        type="pil",
                    )

                    output_seed = gr.Number(
                        label="Seed used",
                        precision=0,
                    )

            generate_btn.click(
                fn=generate_image,
                inputs=[
                    prompt,
                    height,
                    width,
                    steps,
                    seed,
                    randomize,
                ],
                outputs=[
                    output_image,
                    output_seed,
                ],
                api_name="generate_image",
            )

        # ====================================================
        # IMAGE WATERMARK
        # ====================================================

        with gr.Tab("🧹 图片去水印"):

            with gr.Row():

                with gr.Column():

                    wm_input = gr.Image(
                        type="pil",
                        label="Input image",
                    )

                    wm_prompt = gr.Textbox(
                        value="watermark",
                        label="Detection prompt",
                    )

                    wm_max = gr.Slider(
                        minimum=1,
                        maximum=90,
                        value=30,
                        step=1,
                        label="Max bbox area (%)",
                    )

                    wm_transparent = gr.Checkbox(
                        value=False,
                        label="Transparent output",
                    )

                    with gr.Row():

                        detect_btn = gr.Button(
                            "Detect / Preview"
                        )

                        remove_btn = gr.Button(
                            "Remove watermark",
                            variant="primary",
                        )

                with gr.Column():

                    wm_preview_out = gr.Image(
                        label="Detection preview",
                        type="pil",
                    )

                    wm_result = gr.Image(
                        label="Result",
                        type="pil",
                    )

                    wm_mask = gr.Image(
                        label="Mask",
                        type="pil",
                    )

            detect_btn.click(
                fn=wm_preview,
                inputs=[
                    wm_input,
                    wm_prompt,
                    wm_max,
                ],
                outputs=wm_preview_out,
                api_name="watermark_preview",
            )

            remove_btn.click(
                fn=wm_remove_image,
                inputs=[
                    wm_input,
                    wm_prompt,
                    wm_max,
                    wm_transparent,
                ],
                outputs=[
                    wm_result,
                    wm_mask,
                ],
                api_name="remove_image",
            )

        # ====================================================
        # VIDEO WATERMARK
        # ====================================================

        with gr.Tab("🎬 视频去水印"):

            gr.Markdown(
                """
### 视频去水印
ZeroGPU 版本：
- 默认最多 30 秒
- 每 N 帧进行一次 Florence-2 检测
- LaMa 对检测到的水印区域进行局部修复
- 尽量保留原视频音频
"""
            )

            # ------------------------------------------------
            # IMPORTANT:
            # Gradio 5.23.1 does NOT use type="filepath"
            # ------------------------------------------------

            video_input = gr.Video(
                label="Input video (max 30s)",
                sources=["upload"],
                format="mp4",
            )

            video_prompt = gr.Textbox(
                value="watermark",
                label="Detection prompt",
            )

            with gr.Row():

                video_max = gr.Slider(
                    minimum=1,
                    maximum=90,
                    value=30,
                    step=1,
                    label="Max bbox area (%)",
                )

                video_skip = gr.Slider(
                    minimum=1,
                    maximum=30,
                    value=15,
                    step=1,
                    label="Detection interval",
                )

            with gr.Row():

                fade_in = gr.Slider(
                    minimum=0,
                    maximum=30,
                    value=2,
                    step=1,
                    label="Fade in frames",
                )

                fade_out = gr.Slider(
                    minimum=0,
                    maximum=30,
                    value=2,
                    step=1,
                    label="Fade out frames",
                )

            video_btn = gr.Button(
                "Remove watermark from video",
                variant="primary",
            )

            video_output = gr.Video(
                label="Output video",
                format="mp4",
            )

            video_btn.click(
                fn=wm_remove_video,
                inputs=[
                    video_input,
                    video_prompt,
                    video_max,
                    video_skip,
                    fade_in,
                    fade_out,
                ],
                outputs=video_output,
                api_name="remove_video",
            )


# ============================================================
# START
# ============================================================

if __name__ == "__main__":

    print(
        "Starting Gradio Space..."
    )

    print(
        f"CUDA available: "
        f"{torch.cuda.is_available()}"
    )

    if torch.cuda.is_available():

        print(
            "CUDA device:",
            torch.cuda.get_device_name(
                torch.cuda.current_device()
            ),
        )

    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
    )

	import os
	import torch
	import gradio as gr
	import spaces
	import diffusers

	from diffusers import ZImagePipeline


	# ============================================================
	# VERSION CHECK
	# ============================================================

	print("=" * 60)
	print("VERSION CHECK")
	print("=" * 60)
	print("Torch:", torch.__version__)
	print("Diffusers:", diffusers.__version__)
	print(
	"Transformers:",
	__import__("transformers").__version__,
	)
	print(
	"HuggingFace Hub:",
	__import__("huggingface_hub").__version__,
	)
	print(
	"Gradio:",
	gr.__version__,
	)
	print(
	"CUDA available:",
	torch.cuda.is_available(),
	)

	print(
	"Has ZImagePipeline:",
	hasattr(diffusers, "ZImagePipeline"),
	)

	print("=" * 60)


	# ============================================================
	# CHECK Z-IMAGE PIPELINE
	# ============================================================

	if not hasattr(diffusers, "ZImagePipeline"):
	raise RuntimeError(
	"ZImagePipeline is not available. "
	"Please install diffusers==0.36.0 or newer."
	)


	# ============================================================
	# LOAD Z-IMAGE
	# ============================================================

	print("Loading Z-Image-Turbo pipeline...")

	pipe = ZImagePipeline.from_pretrained(
	"Tongyi-MAI/Z-Image-Turbo",
	torch_dtype=torch.bfloat16,
	low_cpu_mem_usage=False,
	)

	pipe.to("cuda")

	print("Z-Image pipeline loaded!")


	# ============================================================
	# WATERMARK MODULE
	# ============================================================

	from watermark import (
	preview,
	remove_image,
	process_video,
	)

	from watermark.video import get_video_info


	# ============================================================
	# Z-IMAGE GENERATION
	# ============================================================

	@spaces.GPU(duration=60)
	def generate_image(
	prompt: str,
	height: int,
	width: int,
	steps: int,
	seed: int,
	randomize_seed: bool,
	progress=gr.Progress(track_tqdm=True),
	):
	"""
	Generate an image with Z-Image-Turbo.
	"""

	if not prompt or not prompt.strip():
	raise gr.Error("Please enter a prompt.")

	if randomize_seed:
	seed = torch.randint(
	0,
	2**32 - 1,
	(1,),
	).item()

	seed = int(seed)

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

	print(
	f"Generating image: "
	f"{width}x{height}, "
	f"steps={steps}, "
	f"seed={seed}"
	)

	image = pipe(
	prompt=prompt,
	height=int(height),
	width=int(width),
	num_inference_steps=int(steps),
	guidance_scale=0.0,
	generator=generator,
	).images[0]

	return image, seed


	# ============================================================
	# WATERMARK PREVIEW
	# ============================================================

	@spaces.GPU(duration=30)
	def wm_preview(
	image,
	detection_prompt: str = "watermark",
	max_bbox_percent: float = 30,
	):
	"""
	Detect likely watermark regions.
	"""

	if image is None:
	raise gr.Error("Please upload an image.")

	return preview(
	image,
	float(max_bbox_percent),
	detection_prompt or "watermark",
	)


	# ============================================================
	# WATERMARK IMAGE REMOVAL
	# ============================================================

	@spaces.GPU(duration=45)
	def wm_remove_image(
	image,
	detection_prompt: str = "watermark",
	max_bbox_percent: float = 30,
	transparent: bool = False,
	):
	"""
	Detect and remove watermark regions from an image.
	"""

	if image is None:
	raise gr.Error("Please upload an image.")

	return remove_image(
	image,
	float(max_bbox_percent),
	detection_prompt or "watermark",
	bool(transparent),
	)


	# ============================================================
	# WATERMARK VIDEO REMOVAL
	# ============================================================

	@spaces.GPU(duration=60)
	def wm_remove_video(
	video_path: str,
	detection_prompt: str = "watermark",
	max_bbox_percent: float = 30,
	detection_skip: int = 5,
	fade_in: int = 2,
	fade_out: int = 2,
	progress=gr.Progress(),
	):
	"""
	Detect and remove watermarks from a video.
	"""

	if not video_path:
	raise gr.Error(
	"Please upload a video."
	)

	# --------------------------------------------------------
	# Inspect video
	# --------------------------------------------------------

	try:

	duration, _, _, _, _ = get_video_info(
	video_path
	)

	max_seconds = float(
	os.getenv(
	"WATERMARK_MAX_VIDEO_SECONDS",
	"30",
	)
	)

	if duration > max_seconds + 0.25:

	raise gr.Error(
	f"Video is {duration:.1f}s. "
	f"Maximum is "
	f"{max_seconds:.0f}s "
	f"for this ZeroGPU build. "
	f"Please trim the video."
	)

	except gr.Error:
	raise

	except Exception as e:

	raise gr.Error(
	f"Unable to inspect video: {e}"
	)

	# --------------------------------------------------------
	# Process video
	# --------------------------------------------------------

	return process_video(
	video_path,
	detection_prompt or "watermark",
	float(max_bbox_percent),
	max(1, int(detection_skip)),
	max(0, int(fade_in)),
	max(0, int(fade_out)),
	progress=progress,
	)


	# ============================================================
	# GRADIO UI
	# ============================================================

	with gr.Blocks(
	title="Z-Image + Watermark Remover V5.2"
	) as demo:

	gr.Markdown(
	"""
	# 🎨 Z-Image + Watermark Remover V5.2

	Z-Image-Turbo + Florence-2 + LaMa

	支持：

	- 🎨 AI 文生图
	- 🧹 图片去水印
	- 🎬 短视频去水印
	- 🔌 Gradio API
	"""
	)

	# ========================================================
	# TABS
	# ========================================================

	with gr.Tabs():

	# ====================================================
	# Z-IMAGE
	# ====================================================

	with gr.Tab("🎨 Z-Image"):

	with gr.Row():

	with gr.Column():

	prompt = gr.Textbox(
	label="Prompt",
	placeholder=(
	"Describe the image..."
	),
	lines=5,
	)

	with gr.Row():

	height = gr.Slider(
	minimum=512,
	maximum=2048,
	value=1024,
	step=64,
	label="Height",
	)

	width = gr.Slider(
	minimum=512,
	maximum=2048,
	value=1024,
	step=64,
	label="Width",
	)

	with gr.Row():

	steps = gr.Slider(
	minimum=1,
	maximum=30,
	value=9,
	step=1,
	label="Inference Steps",
	)

	seed = gr.Number(
	value=42,
	precision=0,
	label="Seed",
	)

	randomize = gr.Checkbox(
	value=True,
	label="Randomize seed",
	)

	generate_btn = gr.Button(
	"Generate",
	variant="primary",
	)

	with gr.Column():

	output_image = gr.Image(
	label="Generated Image",
	type="pil",
	)

	output_seed = gr.Number(
	label="Seed used",
	precision=0,
	)

	generate_btn.click(
	fn=generate_image,
	inputs=[
	prompt,
	height,
	width,
	steps,
	seed,
	randomize,
	],
	outputs=[
	output_image,
	output_seed,
	],
	api_name="generate_image",
	)

	# ====================================================
	# IMAGE WATERMARK
	# ====================================================

	with gr.Tab("🧹 图片去水印"):

	with gr.Row():

	with gr.Column():

	wm_input = gr.Image(
	type="pil",
	label="Input image",
	)

	wm_prompt = gr.Textbox(
	value="watermark",
	label="Detection prompt",
	)

	wm_max = gr.Slider(
	minimum=1,
	maximum=90,
	value=30,
	step=1,
	label="Max bbox area (%)",
	)

	wm_transparent = gr.Checkbox(
	value=False,
	label="Transparent output",
	)

	with gr.Row():

	detect_btn = gr.Button(
	"Detect / Preview"
	)

	remove_btn = gr.Button(
	"Remove watermark",
	variant="primary",
	)

	with gr.Column():

	wm_preview_out = gr.Image(
	label="Detection preview",
	type="pil",
	)

	wm_result = gr.Image(
	label="Result",
	type="pil",
	)

	wm_mask = gr.Image(
	label="Mask",
	type="pil",
	)

	detect_btn.click(
	fn=wm_preview,
	inputs=[
	wm_input,
	wm_prompt,
	wm_max,
	],
	outputs=wm_preview_out,
	api_name="watermark_preview",
	)

	remove_btn.click(
	fn=wm_remove_image,
	inputs=[
	wm_input,
	wm_prompt,
	wm_max,
	wm_transparent,
	],
	outputs=[
	wm_result,
	wm_mask,
	],
	api_name="remove_image",
	)

	# ====================================================
	# VIDEO WATERMARK
	# ====================================================

	with gr.Tab("🎬 视频去水印"):

	gr.Markdown(
	"""
	### 视频去水印

	ZeroGPU 版本：

	- 默认最多 30 秒
	- 每 N 帧进行一次 Florence-2 检测
	- LaMa 对检测到的水印区域进行局部修复
	- 尽量保留原视频音频
	"""
	)

	# ------------------------------------------------
	# IMPORTANT:
	# Gradio 5.23.1 does NOT use type="filepath"
	# ------------------------------------------------

	video_input = gr.Video(
	label="Input video (max 30s)",
	sources=["upload"],
	format="mp4",
	)

	video_prompt = gr.Textbox(
	value="watermark",
	label="Detection prompt",
	)

	with gr.Row():

	video_max = gr.Slider(
	minimum=1,
	maximum=90,
	value=30,
	step=1,
	label="Max bbox area (%)",
	)

	video_skip = gr.Slider(
	minimum=1,
	maximum=30,
	value=15,
	step=1,
	label="Detection interval",
	)

	with gr.Row():

	fade_in = gr.Slider(
	minimum=0,
	maximum=30,
	value=2,
	step=1,
	label="Fade in frames",
	)

	fade_out = gr.Slider(
	minimum=0,
	maximum=30,
	value=2,
	step=1,
	label="Fade out frames",
	)

	video_btn = gr.Button(
	"Remove watermark from video",
	variant="primary",
	)

	video_output = gr.Video(
	label="Output video",
	format="mp4",
	)

	video_btn.click(
	fn=wm_remove_video,
	inputs=[
	video_input,
	video_prompt,
	video_max,
	video_skip,
	fade_in,
	fade_out,
	],
	outputs=video_output,
	api_name="remove_video",
	)


	# ============================================================
	# START
	# ============================================================

	if __name__ == "__main__":

	print(
	"Starting Gradio Space..."
	)

	print(
	f"CUDA available: "
	f"{torch.cuda.is_available()}"
	)

	if torch.cuda.is_available():

	print(
	"CUDA device:",
	torch.cuda.get_device_name(
	torch.cuda.current_device()
	),
	)

	demo.launch(
	server_name="0.0.0.0",
	server_port=7860,
	)