Spaces:
Runtime error
Runtime error
File size: 4,280 Bytes
167dae1 8d3e43b 167dae1 8d3e43b 167dae1 8d3e43b 167dae1 8d3e43b 167dae1 8d3e43b 167dae1 8d3e43b 167dae1 8d3e43b 167dae1 8d3e43b 167dae1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 | import torch
import gradio as gr
import random
from diffusers import DiffusionPipeline
from sdnq import SDNQConfig
from sdnq.loader import apply_sdnq_options_to_model
# --- Model Configuration ---
MODEL_ID = "Abrahamm3r/Z-Image-SDNQ-uint4-svd-r32"
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Loading model: {MODEL_ID} on {device}...")
# 1. Load the pipeline with trust_remote_code=True for Z-Image architecture
# We use bfloat16 as it is standard for these newer flux/z-image models
pipe = DiffusionPipeline.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
trust_remote_code=True
)
# 2. Apply SDNQ quantization hooks to the transformer
# This is critical for the model to run with the compressed weights
pipe.transformer = apply_sdnq_options_to_model(pipe.transformer)
# 3. Optimize memory
if device == "cuda":
pipe.to(device)
# Enable if you are on a smaller GPU (e.g., T4 16GB) to save VRAM
# pipe.enable_model_cpu_offload()
print("Model loaded successfully!")
# --- Helper Functions ---
# Preset resolutions for Aspect Ratios
# Z-Image handles various resolutions, but these are safe standard presets
ASPECT_RATIOS = {
"1:1 (Square)": (1024, 1024),
"16:9 (Cinematic)": (1280, 720),
"9:16 (Portrait)": (720, 1280),
"4:3 (Photo)": (1152, 864),
"3:4 (Portrait Photo)": (864, 1152),
"21:9 (Ultrawide)": (1536, 640)
}
def generate_image(prompt, negative_prompt, steps, aspect_ratio_choice, seed, guidance_scale):
# Determine Width/Height from preset
width, height = ASPECT_RATIOS.get(aspect_ratio_choice, (1024, 1024))
# Handle Seed
if seed == -1:
seed = random.randint(0, 2**32 - 1)
generator = torch.Generator(device=device).manual_seed(int(seed))
print(f"Generating: '{prompt}' | Steps: {steps} | Size: {width}x{height} | Seed: {seed}")
try:
image = pipe(
prompt=prompt,
negative_prompt=negative_prompt,
width=width,
height=height,
num_inference_steps=steps,
guidance_scale=guidance_scale,
generator=generator
).images[0]
return image, seed
except Exception as e:
raise gr.Error(f"Generation failed: {str(e)}")
# --- Gradio UI ---
custom_css = """
#col-container { max-width: 800px; margin: 0 auto; }
#generate-btn { background: linear-gradient(90deg, #4B79A1 0%, #283E51 100%); border: none; color: white; }
"""
with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
with gr.Column(elem_id="col-container"):
gr.Markdown(f"# ⚡ Z-Image SDNQ (uint4-svd-r32) Generator")
gr.Markdown(f"Running `{MODEL_ID}`. This uses Structured Decomposable Neural Quantization for high efficiency.")
with gr.Row():
with gr.Column(scale=2):
prompt = gr.Textbox(label="Prompt", placeholder="Describe the image you want...", lines=3)
negative_prompt = gr.Textbox(label="Negative Prompt", placeholder="Low quality, blurry, ugly...", value="low quality, bad anatomy, worst quality, distortion, blurry")
with gr.Row():
aspect_ratio = gr.Dropdown(
label="Aspect Ratio",
choices=list(ASPECT_RATIOS.keys()),
value="1:1 (Square)"
)
steps = gr.Slider(label="Inference Steps", minimum=4, maximum=50, step=1, value=25)
with gr.Row():
guidance = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=3.5)
seed = gr.Number(label="Seed (-1 for Random)", value=-1, precision=0)
run_btn = gr.Button("Generate Image", elem_id="generate-btn", size="lg")
with gr.Column(scale=2):
result_image = gr.Image(label="Generated Image", type="pil")
seed_output = gr.Label(label="Used Seed")
run_btn.click(
fn=generate_image,
inputs=[prompt, negative_prompt, steps, aspect_ratio, seed, guidance],
outputs=[result_image, seed_output]
)
if __name__ == "__main__":
demo.queue().launch() |