Update app.py

app.py

@@ -1,568 +1,358 @@
-import spaces
-from dataclasses import dataclass
-import json
-import logging
+# ============================================================
+# IMPORTANT: imports order matters for Hugging Face Spaces
+# ============================================================
+
 import os
+import gc
 import random
-import re
-import sys
 import warnings
+import logging
+
+# ---- Spaces GPU decorator (must be imported early) ----------
+try:
+    import spaces  # noqa: F401
+    SPACES_AVAILABLE = True
+except Exception:
+    SPACES_AVAILABLE = False
 
-from PIL import Image
-from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
 import gradio as gr
+import numpy as np
+from PIL import Image
+
 import torch
+from huggingface_hub import login
+
+from diffusers import (
+    ZImagePipeline,
+    ZImageImg2ImgPipeline,
+    AutoencoderKL,
+    FlowMatchEulerDiscreteScheduler,
+)
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from prompt_check import is_unsafe_prompt
-
-sys.path.append(os.path.dirname(os.path.abspath(__file__)))
-
-from diffusers import ZImagePipeline
-from diffusers.models.transformers.transformer_z_image import ZImageTransformer2DModel
+# ============================================================
+# Config
+# ============================================================
 
-from pe import prompt_template
-
-# ==================== Environment Variables ==================================
-MODEL_PATH = os.environ.get("MODEL_PATH", "telcom/dee-z-image")
-ENABLE_COMPILE = os.environ.get("ENABLE_COMPILE", "true").lower() == "true"
-ENABLE_WARMUP = os.environ.get("ENABLE_WARMUP", "true").lower() == "true"
-ATTENTION_BACKEND = os.environ.get("ATTENTION_BACKEND", "flash_3")
-UNSAFE_MAX_NEW_TOKEN = int(os.environ.get("UNSAFE_MAX_NEW_TOKEN", "10"))
-DASHSCOPE_API_KEY = os.environ.get("DASHSCOPE_API_KEY")
-HF_TOKEN = os.environ.get("HF_TOKEN")
-UNSAFE_PROMPT_CHECK = os.environ.get("UNSAFE_PROMPT_CHECK")
-# =============================================================================
+MODEL_PATH = os.environ.get("MODEL_PATH", "telcom/dee-z-image").strip()
+ATTENTION_BACKEND = os.environ.get("ATTENTION_BACKEND", "flash_3").strip()  # try: flash_3, flash, sdpa
+ENABLE_COMPILE = os.environ.get("ENABLE_COMPILE", "false").lower() == "true"
 
+HF_TOKEN = os.getenv("HF_TOKEN", "").strip()
+if HF_TOKEN:
+    login(token=HF_TOKEN)
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 warnings.filterwarnings("ignore")
 logging.getLogger("transformers").setLevel(logging.ERROR)
 
-RES_CHOICES = {
-    "1024": [
-        "1024x1024 ( 1:1 )",
-        "1152x896 ( 9:7 )",
-        "896x1152 ( 7:9 )",
-        "1152x864 ( 4:3 )",
-        "864x1152 ( 3:4 )",
-        "1248x832 ( 3:2 )",
-        "832x1248 ( 2:3 )",
-        "1280x720 ( 16:9 )",
-        "720x1280 ( 9:16 )",
-        "1344x576 ( 21:9 )",
-        "576x1344 ( 9:21 )",
-    ],
-    "1280": [
-        "1280x1280 ( 1:1 )",
-        "1440x1120 ( 9:7 )",
-        "1120x1440 ( 7:9 )",
-        "1472x1104 ( 4:3 )",
-        "1104x1472 ( 3:4 )",
-        "1536x1024 ( 3:2 )",
-        "1024x1536 ( 2:3 )",
-        "1536x864 ( 16:9 )",
-        "864x1536 ( 9:16 )",
-        "1680x720 ( 21:9 )",
-        "720x1680 ( 9:21 )",
-    ],
-    "1536": [
-        "1536x1536 ( 1:1 )",
-        "1728x1344 ( 9:7 )",
-        "1344x1728 ( 7:9 )",
-        "1728x1296 ( 4:3 )",
-        "1296x1728 ( 3:4 )",
-        "1872x1248 ( 3:2 )",
-        "1248x1872 ( 2:3 )",
-        "2048x1152 ( 16:9 )",
-        "1152x2048 ( 9:16 )",
-        "2016x864 ( 21:9 )",
-        "864x2016 ( 9:21 )",
-    ],
-}
+MAX_SEED = np.iinfo(np.int32).max
 
-RESOLUTION_SET = []
-for resolutions in RES_CHOICES.values():
-    RESOLUTION_SET.extend(resolutions)
+# ============================================================
+# Device & dtype
+# ============================================================
 
-EXAMPLE_PROMPTS = [
-    [
-        "Young Chinese woman in red Hanfu, intricate embroidery. Impeccable makeup, red floral forehead pattern. Elaborate high bun, golden phoenix headdress, red flowers, beads. Holds round folding fan with lady, trees, bird. Neon lightning-bolt lamp (⚡️), bright yellow glow, above extended left palm. Soft-lit outdoor night background, silhouetted tiered pagoda (西安大雁塔), blurred colorful distant lights."
-    ],
-    [
-        '''A vertical digital illustration depicting a serene and majestic Chinese landscape, rendered in a style reminiscent of traditional Shanshui painting but with a modern, clean aesthetic. The scene is dominated by towering, steep cliffs in various shades of blue and teal, which frame a central valley. In the distance, layers of mountains fade into a light blue and white mist, creating a strong sense of atmospheric perspective and depth. A calm, turquoise river flows through the center of the composition, with a small, traditional Chinese boat, possibly a sampan, navigating its waters. The boat has a bright yellow canopy and a red hull, and it leaves a gentle wake behind it. It carries several indistinct figures of people. Sparse vegetation, including green trees and some bare-branched trees, clings to the rocky ledges and peaks. The overall lighting is soft and diffused, casting a tranquil glow over the entire scene. Centered in the image is overlaid text. At the top of the text block is a small, red, circular seal-like logo containing stylized characters. Below it, in a smaller, black, sans-serif font, are the words 'Zao-Xiang * East Beauty & West Fashion * Z-Image'. Directly beneath this, in a larger, elegant black serif font, is the word 'SHOW & SHARE CREATIVITY WITH THE WORLD'. Among them, there are "SHOW & SHARE", "CREATIVITY", and "WITH THE WORLD"'''
-    ],
-  
-]
+cuda_available = torch.cuda.is_available()
+device = torch.device("cuda" if cuda_available else "cpu")
 
+if cuda_available and hasattr(torch.cuda, "is_bf16_supported") and torch.cuda.is_bf16_supported():
+    dtype = torch.bfloat16
+elif cuda_available:
+    dtype = torch.float16
+else:
+    dtype = torch.float32
 
-def get_resolution(resolution):
-    match = re.search(r"(\d+)\s*[×x]\s*(\d+)", resolution)
-    if match:
-        return int(match.group(1)), int(match.group(2))
-    return 1024, 1024
+# A conservative max for most Spaces GPUs. Increase if you know you have headroom.
+MAX_IMAGE_SIZE = 1536 if cuda_available else 768
 
+fallback_msg = ""
+if not cuda_available:
+    fallback_msg = "GPU unavailable. Running in CPU fallback mode (slow)."
 
-def load_models(model_path, enable_compile=False, attention_backend="native"):
-    print(f"Loading models from {model_path}...")
+# ============================================================
+# Load pipelines
+# ============================================================
 
-    use_auth_token = HF_TOKEN if HF_TOKEN else True
+pipe_txt2img = None
+pipe_img2img = None
+model_loaded = False
+load_error = None
 
-    if not os.path.exists(model_path):
-        vae = AutoencoderKL.from_pretrained(
-            f"{model_path}",
-            subfolder="vae",
-            torch_dtype=torch.bfloat16,
-            device_map="cuda",
-            use_auth_token=use_auth_token,
-        )
-
-        text_encoder = AutoModelForCausalLM.from_pretrained(
-            f"{model_path}",
-            subfolder="text_encoder",
-            torch_dtype=torch.bfloat16,
-            device_map="cuda",
-            use_auth_token=use_auth_token,
-        ).eval()
-
-        tokenizer = AutoTokenizer.from_pretrained(f"{model_path}", subfolder="tokenizer", use_auth_token=use_auth_token)
-    else:
-        vae = AutoencoderKL.from_pretrained(
-            os.path.join(model_path, "vae"), torch_dtype=torch.bfloat16, device_map="cuda"
-        )
-
-        text_encoder = AutoModelForCausalLM.from_pretrained(
-            os.path.join(model_path, "text_encoder"),
-            torch_dtype=torch.bfloat16,
-            device_map="cuda",
-        ).eval()
-
-        tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_path, "tokenizer"))
+def _try_load_with_from_pretrained():
+    """
+    Preferred path: load everything via Diffusers from_pretrained.
+    Works when the repo is structured as a standard Diffusers pipeline repo.
+    """
+    kwargs = {
+        "torch_dtype": dtype,
+        "use_safetensors": True,
+    }
+    if HF_TOKEN:
+        kwargs["token"] = HF_TOKEN
+
+    p_txt = ZImagePipeline.from_pretrained(MODEL_PATH, **kwargs)
+    p_img = ZImageImg2ImgPipeline(**p_txt.components)
+    return p_txt, p_img
+
+def _fallback_manual_load():
+    """
+    Fallback path: load subfolders manually, similar to many Z-Image demos.
+    Works when MODEL_PATH points to a repo with subfolders:
+      vae/, transformer/, text_encoder/, tokenizer/
+    """
+    use_auth_token = HF_TOKEN if HF_TOKEN else True
 
+    vae = AutoencoderKL.from_pretrained(
+        MODEL_PATH,
+        subfolder="vae",
+        torch_dtype=dtype,
+        use_auth_token=use_auth_token,
+    )
+    text_encoder = AutoModelForCausalLM.from_pretrained(
+        MODEL_PATH,
+        subfolder="text_encoder",
+        torch_dtype=dtype,
+        use_auth_token=use_auth_token,
+    ).eval()
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_PATH,
+        subfolder="tokenizer",
+        use_auth_token=use_auth_token,
+    )
     tokenizer.padding_side = "left"
 
-    if enable_compile:
-        print("Enabling torch.compile optimizations...")
-        torch._inductor.config.conv_1x1_as_mm = True
-        torch._inductor.config.coordinate_descent_tuning = True
-        torch._inductor.config.epilogue_fusion = False
-        torch._inductor.config.coordinate_descent_check_all_directions = True
-        torch._inductor.config.max_autotune_gemm = True
-        torch._inductor.config.max_autotune_gemm_backends = "TRITON,ATEN"
-        torch._inductor.config.triton.cudagraphs = False
-
-    pipe = ZImagePipeline(scheduler=None, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, transformer=None)
-
-    if enable_compile:
-        pipe.vae.disable_tiling()
-
-    if not os.path.exists(model_path):
-        transformer = ZImageTransformer2DModel.from_pretrained(
-            f"{model_path}", subfolder="transformer", use_auth_token=use_auth_token
-        ).to("cuda", torch.bfloat16)
-    else:
-        transformer = ZImageTransformer2DModel.from_pretrained(os.path.join(model_path, "transformer")).to(
-            "cuda", torch.bfloat16
-        )
-
-    pipe.transformer = transformer
-    pipe.transformer.set_attention_backend(attention_backend)
-
-    if enable_compile:
-        print("Compiling transformer...")
-        pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune-no-cudagraphs", fullgraph=False)
-
-    pipe.to("cuda", torch.bfloat16)
-
-    from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
-    from transformers import CLIPImageProcessor
-
-    safety_model_id = "CompVis/stable-diffusion-safety-checker"
-    safety_feature_extractor = CLIPImageProcessor.from_pretrained(safety_model_id)
-    safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id, torch_dtype=torch.float16).to("cuda")
-
-    pipe.safety_feature_extractor = safety_feature_extractor
-    pipe.safety_checker = safety_checker
-    return pipe
-
-
-def generate_image(
-    pipe,
-    prompt,
-    resolution="1024x1024",
-    seed=42,
-    guidance_scale=5.0,
-    num_inference_steps=50,
-    shift=3.0,
-    max_sequence_length=512,
-    progress=gr.Progress(track_tqdm=True),
-):
-    width, height = get_resolution(resolution)
-
-    generator = torch.Generator("cuda").manual_seed(seed)
-
-    scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=shift)
-    pipe.scheduler = scheduler
-
-    image = pipe(
-        prompt=prompt,
-        height=height,
-        width=width,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        generator=generator,
-        max_sequence_length=max_sequence_length,
-    ).images[0]
-
-    return image
-
-
-def warmup_model(pipe, resolutions):
-    print("Starting warmup phase...")
-
-    dummy_prompt = "warmup"
-
-    for res_str in resolutions:
-        print(f"Warming up for resolution: {res_str}")
-        try:
-            for i in range(3):
-                generate_image(
-                    pipe,
-                    prompt=dummy_prompt,
-                    resolution=res_str,
-                    num_inference_steps=9,
-                    guidance_scale=0.0,
-                    seed=42 + i,
-                )
-        except Exception as e:
-            print(f"Warmup failed for {res_str}: {e}")
-
-    print("Warmup completed.")
-
-
-# ==================== Prompt Expander ====================
-@dataclass
-class PromptOutput:
-    status: bool
-    prompt: str
-    seed: int
-    system_prompt: str
-    message: str
-
-
-class PromptExpander:
-    def __init__(self, backend="api", **kwargs):
-        self.backend = backend
-
-    def decide_system_prompt(self, template_name=None):
-        return prompt_template
-
-
-class APIPromptExpander(PromptExpander):
-    def __init__(self, api_config=None, **kwargs):
-        super().__init__(backend="api", **kwargs)
-        self.api_config = api_config or {}
-        self.client = self._init_api_client()
-
-    def _init_api_client(self):
-        try:
-            from openai import OpenAI
-
-            api_key = self.api_config.get("api_key") or DASHSCOPE_API_KEY
-            base_url = self.api_config.get("base_url", "https://dashscope.aliyuncs.com/compatible-mode/v1")
-
-            if not api_key:
-                print("Warning: DASHSCOPE_API_KEY not found.")
-                return None
-
-            return OpenAI(api_key=api_key, base_url=base_url)
-        except ImportError:
-            print("Please install openai: pip install openai")
-            return None
-        except Exception as e:
-            print(f"Failed to initialize API client: {e}")
-            return None
-
-    def __call__(self, prompt, system_prompt=None, seed=-1, **kwargs):
-        return self.extend(prompt, system_prompt, seed, **kwargs)
-
-    def extend(self, prompt, system_prompt=None, seed=-1, **kwargs):
-        if self.client is None:
-            return PromptOutput(False, "", seed, system_prompt, "API client not initialized")
-
-        if system_prompt is None:
-            system_prompt = self.decide_system_prompt()
-
-        if "{prompt}" in system_prompt:
-            system_prompt = system_prompt.format(prompt=prompt)
-            prompt = " "
-
-        try:
-            model = self.api_config.get("model", "qwen3-max-preview")
-            response = self.client.chat.completions.create(
-                model=model,
-                messages=[{"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}],
-                temperature=0.7,
-                top_p=0.8,
-            )
-
-            content = response.choices[0].message.content
-            json_start = content.find("```json")
-            if json_start != -1:
-                json_end = content.find("```", json_start + 7)
-                try:
-                    json_str = content[json_start + 7 : json_end].strip()
-                    data = json.loads(json_str)
-                    expanded_prompt = data.get("revised_prompt", content)
-                except:
-                    expanded_prompt = content
-            else:
-                expanded_prompt = content
-
-            return PromptOutput(
-                status=True, prompt=expanded_prompt, seed=seed, system_prompt=system_prompt, message=content
-            )
-        except Exception as e:
-            return PromptOutput(False, "", seed, system_prompt, str(e))
-
-
-def create_prompt_expander(backend="api", **kwargs):
-    if backend == "api":
-        return APIPromptExpander(**kwargs)
-    raise ValueError("Only 'api' backend is supported.")
-
-
-pipe = None
-prompt_expander = None
+    # ZImageTransformer2DModel lives inside diffusers; importing lazily avoids import issues on older versions.
+    from diffusers.models.transformers.transformer_z_image import ZImageTransformer2DModel
 
+    transformer = ZImageTransformer2DModel.from_pretrained(
+        MODEL_PATH,
+        subfolder="transformer",
+        torch_dtype=dtype,
+        use_auth_token=use_auth_token,
+    )
 
-def init_app():
-    global pipe, prompt_expander
+    p_txt = ZImagePipeline(scheduler=None, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, transformer=transformer)
+    p_img = ZImageImg2ImgPipeline(scheduler=None, vae=vae, text_encoder=text_encoder, tokenizer=tokenizer, transformer=transformer)
+    return p_txt, p_img
 
+try:
+    pipe_txt2img, pipe_img2img = _try_load_with_from_pretrained()
+    model_loaded = True
+except Exception as e1:
     try:
-        pipe = load_models(MODEL_PATH, enable_compile=ENABLE_COMPILE, attention_backend=ATTENTION_BACKEND)
-        print(f"Model loaded. Compile: {ENABLE_COMPILE}, Backend: {ATTENTION_BACKEND}")
-
-        if ENABLE_WARMUP:
-            all_resolutions = []
-            for cat in RES_CHOICES.values():
-                all_resolutions.extend(cat)
-            warmup_model(pipe, all_resolutions)
+        pipe_txt2img, pipe_img2img = _fallback_manual_load()
+        model_loaded = True
+    except Exception as e2:
+        load_error = f"from_pretrained error: {repr(e1)}\nmanual_load error: {repr(e2)}"
+        model_loaded = False
 
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        pipe = None
+if model_loaded:
+    pipe_txt2img = pipe_txt2img.to(device)
+    pipe_img2img = pipe_img2img.to(device)
 
+    # Try attention backend (best-effort)
     try:
-        prompt_expander = create_prompt_expander(backend="api", api_config={"model": "qwen3-max-preview"})
-        print("Prompt expander initialized.")
-    except Exception as e:
-        print(f"Error initializing prompt expander: {e}")
-        prompt_expander = None
-
-
-def prompt_enhance(prompt, enable_enhance):
-    if not enable_enhance or not prompt_expander:
-        return prompt, "Enhancement disabled or not available."
-
-    if not prompt.strip():
-        return "", "Please enter a prompt."
-
-    try:
-        result = prompt_expander(prompt)
-        if result.status:
-            return result.prompt, result.message
-        else:
-            return prompt, f"Enhancement failed: {result.message}"
-    except Exception as e:
-        return prompt, f"Error: {str(e)}"
-
-
-@spaces.GPU
-def generate(
-    prompt,
-    resolution="1024x1024 ( 1:1 )",
-    seed=42,
-    steps=9,
-    shift=3.0,
-    random_seed=True,
-    gallery_images=None,
-    enhance=False,
-    progress=gr.Progress(track_tqdm=True),
-):
-    """
-    Generate an image using the Z-Image model based on the provided prompt and settings.
-
-    This function is triggered when the user clicks the "Generate" button. It processes
-    the input prompt (optionally enhancing it), configures generation parameters, and
-    produces an image using the Z-Image diffusion transformer pipeline.
-
-    Args:
-        prompt (str): Text prompt describing the desired image content
-        resolution (str): Output resolution in format "WIDTHxHEIGHT ( RATIO )" (e.g., "1024x1024 ( 1:1 )")
-        seed (int): Seed for reproducible generation
-        steps (int): Number of inference steps for the diffusion process
-        shift (float): Time shift parameter for the flow matching scheduler
-        random_seed (bool): Whether to generate a new random seed, if True will ignore the seed input
-        gallery_images (list): List of previously generated images to append to (only needed for the Gradio UI)
-        enhance (bool): This was Whether to enhance the prompt (DISABLED! Do not use)
-        progress (gr.Progress): Gradio progress tracker for displaying generation progress (only needed for the Gradio UI)
-
-    Returns:
-        tuple: (gallery_images, seed_str, seed_int)
-            - gallery_images: Updated list of generated images including the new image
-            - seed_str: String representation of the seed used for generation
-            - seed_int: Integer representation of the seed used for generation
-    """
-
-    if random_seed:
-        new_seed = random.randint(1, 1000000)
-    else:
-        new_seed = seed if seed != -1 else random.randint(1, 1000000)
-
-    class UnsafeContentError(Exception):
+        if hasattr(pipe_txt2img, "transformer") and hasattr(pipe_txt2img.transformer, "set_attention_backend"):
+            pipe_txt2img.transformer.set_attention_backend(ATTENTION_BACKEND)
+            pipe_img2img.transformer.set_attention_backend(ATTENTION_BACKEND)
+    except Exception:
         pass
 
-    try:
-        if pipe is None:
-            raise gr.Error("Model not loaded.")
-
-        has_unsafe_concept = is_unsafe_prompt(
-            pipe.text_encoder,
-            pipe.tokenizer,
-            system_prompt=UNSAFE_PROMPT_CHECK,
-            user_prompt=prompt,
-            max_new_token=UNSAFE_MAX_NEW_TOKEN,
-        )
-        if has_unsafe_concept:
-            raise UnsafeContentError("Input unsafe")
-
-        final_prompt = prompt
-
-        if enhance:
-            final_prompt, _ = prompt_enhance(prompt, True)
-            print(f"Enhanced prompt: {final_prompt}")
-
+    # Optional compile (best-effort, can break on some setups)
+    if ENABLE_COMPILE and device.type == "cuda":
         try:
-            resolution_str = resolution.split(" ")[0]
-        except:
-            resolution_str = "1024x1024"
-
-        image = generate_image(
-            pipe=pipe,
-            prompt=final_prompt,
-            resolution=resolution_str,
-            seed=new_seed,
-            guidance_scale=0.0,
-            num_inference_steps=int(steps + 1),
-            shift=shift,
-        )
-
-        safety_checker_input = pipe.safety_feature_extractor([image], return_tensors="pt").pixel_values.cuda()
-        _, has_nsfw_concept =  pipe.safety_checker(images=[torch.zeros(1)], clip_input=safety_checker_input)
-        has_nsfw_concept = has_nsfw_concept[0]
-        if has_nsfw_concept:
-            print("input unsafe")
+            pipe_txt2img.transformer = torch.compile(pipe_txt2img.transformer, mode="max-autotune-no-cudagraphs", fullgraph=False)
+            pipe_img2img.transformer = pipe_txt2img.transformer
+        except Exception:
+            pass
 
-    except UnsafeContentError:
-        image = Image.open("nsfw.png")
+    # Disable diffusers progress bars
+    try:
+        pipe_txt2img.set_progress_bar_config(disable=True)
+        pipe_img2img.set_progress_bar_config(disable=True)
+    except Exception:
+        pass
 
-    if gallery_images is None:
-        gallery_images = []
-    # gallery_images.append(image)
-    gallery_images = [image] + gallery_images # latest output to be at the top of the list
+# ============================================================
+# Utility: error image
+# ============================================================
+
+def make_error_image(w, h):
+    return Image.new("RGB", (w, h), (18, 18, 22))
+
+def _prep_init_image(init_image, width, height):
+    if init_image is None:
+        return None
+    if not isinstance(init_image, Image.Image):
+        return None
+    init_image = init_image.convert("RGB")
+    if init_image.size != (width, height):
+        init_image = init_image.resize((width, height), Image.LANCZOS)
+    return init_image
+
+# ============================================================
+# Inference
+# ============================================================
+
+def _infer_impl(
+    prompt: str,
+    negative_prompt: str,
+    seed: int,
+    randomize_seed: bool,
+    width: int,
+    height: int,
+    guidance_scale: float,
+    num_inference_steps: int,
+    shift: float,
+    max_sequence_length: int,
+    init_image,
+    strength: float,
+):
+    width = int(width)
+    height = int(height)
+    seed = int(seed)
 
-    return gallery_images, str(new_seed), int(new_seed)
+    if not model_loaded:
+        return make_error_image(width, height), f"Model load failed:\n\n{load_error}"
 
+    prompt = (prompt or "").strip()
+    if not prompt:
+        return make_error_image(width, height), "Error: Prompt is empty."
 
-init_app()
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
 
-# ==================== AoTI (Ahead of Time Inductor compilation) ====================
+    init_image = _prep_init_image(init_image, width, height)
 
-pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
-spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
+    generator = torch.Generator(device=device)
+    generator = generator.manual_seed(seed)
 
-with gr.Blocks(title="Z-Image Demo") as demo:
-    gr.Markdown(
-        """<div align="center">
+    status = f"Seed: {seed}"
+    if fallback_msg:
+        status += f" | {fallback_msg}"
 
-# Z-Image Generation Demo
+    # Set scheduler per-run because shift can change
+    scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=float(shift))
+    pipe_txt2img.scheduler = scheduler
+    pipe_img2img.scheduler = scheduler
 
-[![GitHub](https://img.shields.io/badge/GitHub-Z--Image-181717?logo=github&logoColor=white)](https://github.com/Tongyi-MAI/Z-Image)
+    try:
+        common_kwargs = dict(
+            prompt=prompt,
+            negative_prompt=(negative_prompt or "").strip() if (guidance_scale and float(guidance_scale) > 1.0) else None,
+            guidance_scale=float(guidance_scale),
+            num_inference_steps=int(num_inference_steps),
+            generator=generator,
+            height=height,
+            width=width,
+            max_sequence_length=int(max_sequence_length),
+        )
 
-*An Efficient Image Generation Foundation Model with Single-Stream Diffusion Transformer*
+        with torch.inference_mode():
+            if device.type == "cuda":
+                with torch.autocast("cuda", dtype=dtype):
+                    if init_image is not None:
+                        out = pipe_img2img(
+                            image=init_image,
+                            strength=float(strength),
+                            **common_kwargs,
+                        )
+                    else:
+                        out = pipe_txt2img(**common_kwargs)
+            else:
+                if init_image is not None:
+                    out = pipe_img2img(
+                        image=init_image,
+                        strength=float(strength),
+                        **common_kwargs,
+                    )
+                else:
+                    out = pipe_txt2img(**common_kwargs)
+
+        image = out.images[0]
+        return image, status
 
-</div>"""
-    )
+    except Exception as e:
+        return make_error_image(width, height), f"Error: {type(e).__name__}: {e}"
+
+    finally:
+        gc.collect()
+        if device.type == "cuda":
+            torch.cuda.empty_cache()
+
+# IMPORTANT: decorator must be explicit
+if SPACES_AVAILABLE:
+    @spaces.GPU
+    def infer(*args, **kwargs):
+        return _infer_impl(*args, **kwargs)
+else:
+    def infer(*args, **kwargs):
+        return _infer_impl(*args, **kwargs)
+
+# ============================================================
+# UI
+# ============================================================
+
+CSS = """
+body {
+    background: #000;
+    color: #fff;
+}
+"""
 
-    with gr.Row():
-        with gr.Column(scale=1):
-            prompt_input = gr.Textbox(label="Prompt", lines=3, placeholder="Enter your prompt here...")
-            # PE components (Temporarily disabled)
-            # with gr.Row():
-            #     enable_enhance = gr.Checkbox(label="Enhance Prompt (DashScope)", value=False)
-            #     enhance_btn = gr.Button("Enhance Only")
-
-            with gr.Row():
-                choices = [int(k) for k in RES_CHOICES.keys()]
-                res_cat = gr.Dropdown(value=1024, choices=choices, label="Resolution Category")
-
-                initial_res_choices = RES_CHOICES["1024"]
-                resolution = gr.Dropdown(
-                    value=initial_res_choices[0], choices=RESOLUTION_SET, label="Width x Height (Ratio)"
-                )
-
-            with gr.Row():
-                seed = gr.Number(label="Seed", value=42, precision=0)
-                random_seed = gr.Checkbox(label="Random Seed", value=True)
-
-            with gr.Row():
-                steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=8, step=1, interactive=False)
-                shift = gr.Slider(label="Time Shift", minimum=1.0, maximum=10.0, value=3.0, step=0.1)
-
-            generate_btn = gr.Button("Generate", variant="primary")
-
-            # Example prompts
-            gr.Markdown("### 📝 Example Prompts")
-            gr.Examples(examples=EXAMPLE_PROMPTS, inputs=prompt_input, label=None)
-
-        with gr.Column(scale=1):
-            output_gallery = gr.Gallery(
-                label="Generated Images",
-                columns=2,
-                rows=2,
-                height=600,
-                object_fit="contain",
-                format="png",
-                interactive=False,
-            )
-            used_seed = gr.Textbox(label="Seed Used", interactive=False)
-
-    def update_res_choices(_res_cat):
-        if str(_res_cat) in RES_CHOICES:
-            res_choices = RES_CHOICES[str(_res_cat)]
-        else:
-            res_choices = RES_CHOICES["1024"]
-        return gr.update(value=res_choices[0], choices=res_choices)
-
-    res_cat.change(update_res_choices, inputs=res_cat, outputs=resolution, api_visibility="private")
-
-    # PE enhancement button (Temporarily disabled)
-    # enhance_btn.click(
-    #     prompt_enhance,
-    #     inputs=[prompt_input, enable_enhance],
-    #     outputs=[prompt_input, final_prompt_output]
-    # )
-
-    generate_btn.click(
-        generate,
-        inputs=[prompt_input, resolution, seed, steps, shift, random_seed, output_gallery],
-        outputs=[output_gallery, used_seed, seed],
-        api_visibility="public",
+with gr.Blocks(title="Z-Image txt2img + img2img") as demo:
+    gr.HTML(f"<style>{CSS}</style>")
+
+    if fallback_msg:
+        gr.Markdown(f"**{fallback_msg}**")
+
+    if not model_loaded:
+        gr.Markdown(f"⚠️ Model failed to load:\n\n{load_error}")
+
+    gr.Markdown("## Z-Image Generator (txt2img + img2img)")
+
+    prompt = gr.Textbox(label="Prompt", lines=2, placeholder="Describe what you want...")
+    init_image = gr.Image(label="Initial image (optional)", type="pil")
+
+    run_button = gr.Button("Generate")
+    result = gr.Image(label="Result")
+    status = gr.Markdown("")
+
+    with gr.Accordion("Advanced Settings", open=False):
+        negative_prompt = gr.Textbox(label="Negative prompt (only used if Guidance > 1)")
+        seed = gr.Slider(0, MAX_SEED, step=1, value=0, label="Seed")
+        randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
+
+        width = gr.Slider(256, MAX_IMAGE_SIZE, step=64, value=1024, label="Width")
+        height = gr.Slider(256, MAX_IMAGE_SIZE, step=64, value=1024, label="Height")
+
+        guidance_scale = gr.Slider(0.0, 10.0, step=0.1, value=0.0, label="Guidance scale")
+        num_inference_steps = gr.Slider(1, 50, step=1, value=8, label="Steps")
+        shift = gr.Slider(1.0, 10.0, step=0.1, value=3.0, label="Time shift")
+
+        max_sequence_length = gr.Slider(64, 512, step=64, value=512, label="Max sequence length")
+
+        strength = gr.Slider(0.0, 1.0, step=0.05, value=0.6, label="Image strength (img2img)")
+
+    run_button.click(
+        fn=infer,
+        inputs=[
+            prompt,
+            negative_prompt,
+            seed,
+            randomize_seed,
+            width,
+            height,
+            guidance_scale,
+            num_inference_steps,
+            shift,
+            max_sequence_length,
+            init_image,
+            strength,
+        ],
+        outputs=[result, status],
     )
 
-css = """
-.fillable{max-width: 1230px !important}
-"""
 if __name__ == "__main__":
-    demo.launch(css=css, mcp_server=True)
+    # Keep the same launch feel as your first script
+    demo.queue().launch(ssr_mode=False)