Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py
 # ============================================================
 # IMPORTANT: imports order matters for Hugging Face Spaces
 # ============================================================
@@ -22,12 +23,20 @@ from PIL import Image
 import torch
 from huggingface_hub import login
 
-from diffusers import (
-    ZImagePipeline,
-    ZImageImg2ImgPipeline,
-    AutoencoderKL,
-    FlowMatchEulerDiscreteScheduler,
-)
+# ---- Diffusers imports (with safe fallbacks) ----------------
+try:
+    from diffusers import ZImagePipeline
+except Exception:
+    # Older/newer diffusers sometimes do not export ZImagePipeline at top-level
+    from diffusers.pipelines.z_image.pipeline_z_image import ZImagePipeline
+
+try:
+    from diffusers import AutoPipelineForImage2Image
+except Exception:
+    # Rare fallback if top-level export is missing
+    from diffusers.pipelines.auto_pipeline import AutoPipelineForImage2Image
+
+from diffusers import AutoencoderKL, FlowMatchEulerDiscreteScheduler
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # ============================================================
@@ -35,7 +44,9 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 # ============================================================
 
 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
+
+# Optional knobs
+ATTENTION_BACKEND = os.environ.get("ATTENTION_BACKEND", "flash_3").strip()  # flash_3, flash, sdpa, native
 ENABLE_COMPILE = os.environ.get("ENABLE_COMPILE", "false").lower() == "true"
 
 HF_TOKEN = os.getenv("HF_TOKEN", "").strip()
@@ -62,7 +73,6 @@ elif cuda_available:
 else:
     dtype = torch.float32
 
-# 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 = ""
@@ -78,199 +88,219 @@ pipe_img2img = None
 model_loaded = False
 load_error = None
 
-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.
+def _manual_load_zimage(model_path: str):
     """
-    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:
+    Manual loader (matches common Z-Image repo layout with subfolders):
       vae/, transformer/, text_encoder/, tokenizer/
+    Works for both remote HF repos and local paths.
     """
     use_auth_token = HF_TOKEN if HF_TOKEN else True
 
+    # Load VAE
     vae = AutoencoderKL.from_pretrained(
-        MODEL_PATH,
+        model_path,
         subfolder="vae",
         torch_dtype=dtype,
         use_auth_token=use_auth_token,
     )
+
+    # Load text encoder + tokenizer
     text_encoder = AutoModelForCausalLM.from_pretrained(
-        MODEL_PATH,
+        model_path,
         subfolder="text_encoder",
         torch_dtype=dtype,
         use_auth_token=use_auth_token,
     ).eval()
+
     tokenizer = AutoTokenizer.from_pretrained(
-        MODEL_PATH,
+        model_path,
         subfolder="tokenizer",
         use_auth_token=use_auth_token,
     )
     tokenizer.padding_side = "left"
 
-    # ZImageTransformer2DModel lives inside diffusers; importing lazily avoids import issues on older versions.
+    # Load transformer
     from diffusers.models.transformers.transformer_z_image import ZImageTransformer2DModel
 
     transformer = ZImageTransformer2DModel.from_pretrained(
-        MODEL_PATH,
+        model_path,
         subfolder="transformer",
         torch_dtype=dtype,
         use_auth_token=use_auth_token,
     )
 
-    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
+    # Build base txt2img pipeline
+    p_txt = ZImagePipeline(
+        scheduler=None,
+        vae=vae,
+        text_encoder=text_encoder,
+        tokenizer=tokenizer,
+        transformer=transformer,
+    )
+    return p_txt
 
 try:
-    pipe_txt2img, pipe_img2img = _try_load_with_from_pretrained()
-    model_loaded = True
-except Exception as e1:
+    # Try standard from_pretrained first (if repo is fully diffusers-compatible)
+    kwargs = {"torch_dtype": dtype, "use_safetensors": True}
+    if HF_TOKEN:
+        kwargs["token"] = HF_TOKEN
+
     try:
-        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
+        pipe_txt2img = ZImagePipeline.from_pretrained(MODEL_PATH, **kwargs)
+    except Exception:
+        pipe_txt2img = _manual_load_zimage(MODEL_PATH)
 
-if model_loaded:
     pipe_txt2img = pipe_txt2img.to(device)
-    pipe_img2img = pipe_img2img.to(device)
 
-    # Try attention backend (best-effort)
+    # Optional attention backend (best-effort)
     try:
         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
 
-    # Optional compile (best-effort, can break on some setups)
+    # Optional compile (best-effort)
     if ENABLE_COMPILE and device.type == "cuda":
         try:
-            pipe_txt2img.transformer = torch.compile(pipe_txt2img.transformer, mode="max-autotune-no-cudagraphs", fullgraph=False)
-            pipe_img2img.transformer = pipe_txt2img.transformer
+            pipe_txt2img.transformer = torch.compile(
+                pipe_txt2img.transformer,
+                mode="max-autotune-no-cudagraphs",
+                fullgraph=False,
+            )
         except Exception:
             pass
 
-    # Disable diffusers progress bars
+    # Disable progress bars
     try:
         pipe_txt2img.set_progress_bar_config(disable=True)
+    except Exception:
+        pass
+
+    # Build img2img pipeline using AutoPipelineForImage2Image
+    # Preferred: from_pipe reuses components without loading a second copy.
+    try:
+        pipe_img2img = AutoPipelineForImage2Image.from_pipe(pipe_txt2img).to(device)
+    except Exception:
+        # Fallback: load an image2image pipeline from the repo (may use more VRAM)
+        pipe_img2img = AutoPipelineForImage2Image.from_pretrained(MODEL_PATH, **kwargs).to(device)
+
+    try:
         pipe_img2img.set_progress_bar_config(disable=True)
     except Exception:
         pass
 
+    model_loaded = True
+
+except Exception as e:
+    load_error = repr(e)
+    model_loaded = False
+
 # ============================================================
-# Utility: error image
+# Utility helpers
 # ============================================================
 
-def make_error_image(w, h):
-    return Image.new("RGB", (w, h), (18, 18, 22))
+def make_error_image(w: int, h: int) -> Image.Image:
+    return Image.new("RGB", (int(w), int(h)), (18, 18, 22))
 
-def _prep_init_image(init_image, width, height):
+def _prep_init_image(init_image, width: int, height: int):
     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
+    img = init_image.convert("RGB")
+    if img.size != (width, height):
+        img = img.resize((width, height), Image.LANCZOS)
+    return img
 
 # ============================================================
 # 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,
+    prompt,
+    negative_prompt,
+    seed,
+    randomize_seed,
+    width,
+    height,
+    guidance_scale,
+    num_inference_steps,
+    shift,
+    max_sequence_length,
     init_image,
-    strength: float,
+    strength,
 ):
     width = int(width)
     height = int(height)
     seed = int(seed)
 
     if not model_loaded:
-        return make_error_image(width, height), f"Model load failed:\n\n{load_error}"
+        return make_error_image(width, height), f"Model load failed: {load_error}"
 
     prompt = (prompt or "").strip()
     if not prompt:
-        return make_error_image(width, height), "Error: Prompt is empty."
+        return make_error_image(width, height), "Error: prompt is empty."
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
 
+    generator = torch.Generator(device=device).manual_seed(seed)
     init_image = _prep_init_image(init_image, width, height)
 
-    generator = torch.Generator(device=device)
-    generator = generator.manual_seed(seed)
-
     status = f"Seed: {seed}"
     if fallback_msg:
         status += f" | {fallback_msg}"
 
-    # Set scheduler per-run because shift can change
+    # Set scheduler per-run so shift changes take effect
     scheduler = FlowMatchEulerDiscreteScheduler(num_train_timesteps=1000, shift=float(shift))
-    pipe_txt2img.scheduler = scheduler
-    pipe_img2img.scheduler = scheduler
+    try:
+        pipe_txt2img.scheduler = scheduler
+    except Exception:
+        pass
+    try:
+        pipe_img2img.scheduler = scheduler
+    except Exception:
+        pass
 
     try:
+        gs = float(guidance_scale)
+        steps = int(num_inference_steps)
+        msl = int(max_sequence_length)
+        st = float(strength)
+
+        # Some pipelines only accept negative_prompt when guidance is used.
+        neg = (negative_prompt or "").strip()
+        if not neg:
+            neg = None
+
         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),
+            guidance_scale=gs,
+            num_inference_steps=steps,
             generator=generator,
             height=height,
             width=width,
-            max_sequence_length=int(max_sequence_length),
+            max_sequence_length=msl,
         )
 
+        if neg is not None:
+            common_kwargs["negative_prompt"] = neg
+
         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,
-                        )
+                        out = pipe_img2img(image=init_image, strength=st, **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,
-                    )
+                    out = pipe_img2img(image=init_image, strength=st, **common_kwargs)
                 else:
                     out = pipe_txt2img(**common_kwargs)
 
-        image = out.images[0]
-        return image, status
+        img = out.images[0]
+        return img, status
 
     except Exception as e:
         return make_error_image(width, height), f"Error: {type(e).__name__}: {e}"
@@ -280,7 +310,7 @@ def _infer_impl(
         if device.type == "cuda":
             torch.cuda.empty_cache()
 
-# IMPORTANT: decorator must be explicit
+# Decorated entrypoint for Spaces
 if SPACES_AVAILABLE:
     @spaces.GPU
     def infer(*args, **kwargs):
@@ -290,7 +320,7 @@ else:
         return _infer_impl(*args, **kwargs)
 
 # ============================================================
-# UI
+# UI (your first style)
 # ============================================================
 
 CSS = """
@@ -319,7 +349,7 @@ with gr.Blocks(title="Z-Image txt2img + img2img") as demo:
     status = gr.Markdown("")
 
     with gr.Accordion("Advanced Settings", open=False):
-        negative_prompt = gr.Textbox(label="Negative prompt (only used if Guidance > 1)")
+        negative_prompt = gr.Textbox(label="Negative prompt (optional)")
         seed = gr.Slider(0, MAX_SEED, step=1, value=0, label="Seed")
         randomize_seed = gr.Checkbox(value=True, label="Randomize seed")
 
@@ -327,9 +357,8 @@ with gr.Blocks(title="Z-Image txt2img + img2img") as demo:
         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")
+        num_inference_steps = gr.Slider(1, 100, 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)")
@@ -354,5 +383,4 @@ with gr.Blocks(title="Z-Image txt2img + img2img") as demo:
     )
 
 if __name__ == "__main__":
-    # Keep the same launch feel as your first script
     demo.queue().launch(ssr_mode=False)