app.py

import spaces
import gradio as gr
import torch
from PIL import Image
import os
import sys
import subprocess
import tempfile
from pathlib import Path
import glob

# Default negative prompts
NEGATIVE_PROMPT_CN = "泛黄，发绿，模糊，低分辨率，低质量图像，扭曲的肢体，诡异的外观，丑陋，AI感，噪点，网格感，JPEG压缩条纹，异常的肢体，水印，乱码，意义不明的字符"
NEGATIVE_PROMPT_EN = "Yellowed, green-tinted, blurry, low-resolution, low-quality image, distorted limbs, eerie appearance, ugly, AI-looking, noise, grid-like artifacts, JPEG compression artifacts, abnormal limbs, watermark, garbled text, meaningless characters"

# Model paths - can be overridden via environment variables
MODELS_DIR = Path(os.environ.get("ZIMAGE_MODELS_DIR", "./models"))


# =============================================================================
# Model Download Functions
# =============================================================================

def download_hf_models(output_dir: Path) -> dict:
    """
    Download required models from Hugging Face using huggingface_hub.
    
    Downloads:
    - DiffSynth-Studio/Z-Image-i2L
    - Tongyi-MAI/Z-Image
    - DiffSynth-Studio/General-Image-Encoders
    - Tongyi-MAI/Z-Image-Turbo  
    
    Returns dict with paths to downloaded models.
    """
    from huggingface_hub import snapshot_download
    
    output_dir.mkdir(parents=True, exist_ok=True)
    
    models = [
        {
            "repo_id": "DiffSynth-Studio/General-Image-Encoders",
            "description": "General Image Encoders (SigLIP2-G384, DINOv3-7B)",
            "allow_patterns": None,
        },
        {
            "repo_id": "Tongyi-MAI/Z-Image-Turbo",
            "description": "Z-Image Turbo (text encoder, VAE, tokenizer)",
            "allow_patterns": [
                "text_encoder/*.safetensors",
                "vae/*.safetensors",
                "tokenizer/*",
            ],
        },
        {
            "repo_id": "Tongyi-MAI/Z-Image",
            "description": "Z-Image base model (transformer)",
            "allow_patterns": ["transformer/*.safetensors"],
        },
        {
            "repo_id": "DiffSynth-Studio/Z-Image-i2L",
            "description": "Z-Image-i2L (Image to LoRA model)",
            "allow_patterns": ["*.safetensors"],
        },
    ]
    
    downloaded_paths = {}
    
    for model in models:
        repo_id = model["repo_id"]
        local_dir = output_dir / repo_id
        
        # Check if already downloaded
        if local_dir.exists() and any(local_dir.rglob("*.safetensors")):
            print(f"   ✓ {repo_id} (already downloaded)")
            downloaded_paths[repo_id] = local_dir
            continue
        
        print(f"   📥 Downloading {repo_id}...")
        print(f"      {model['description']}")
        
        try:
            result_path = snapshot_download(
                repo_id=repo_id,
                local_dir=str(local_dir),
                allow_patterns=model["allow_patterns"],
                local_dir_use_symlinks=False,
                resume_download=True,
            )
            downloaded_paths[repo_id] = Path(result_path)
            print(f"   ✓ {repo_id}")
        except Exception as e:
            print(f"   ❌ Error downloading {repo_id}: {e}")
            raise
    
    return downloaded_paths


def get_model_files(base_path: Path, pattern: str) -> list:
    """Get list of files matching a glob pattern."""
    full_pattern = str(base_path / pattern)
    files = sorted(glob.glob(full_pattern))
    return files


def install_diffsynth_studio():
    """Clone and install DiffSynth-Studio if not already installed."""
    try:
        from diffsynth.pipelines.z_image import ZImagePipeline
        return True, "✅ DiffSynth-Studio is already installed."
    except ImportError:
        pass
    
    repo_dir = Path(__file__).parent / "DiffSynth-Studio"
    
    try:
        if not repo_dir.exists():
            print("📥 Cloning DiffSynth-Studio repository...")
            subprocess.run(
                ["git", "clone", "https://github.com/modelscope/DiffSynth-Studio.git", str(repo_dir)],
                capture_output=True,
                text=True,
                check=True
            )
            print("✅ Repository cloned successfully.")
        else:
            print("📁 DiffSynth-Studio directory already exists, pulling latest...")
            subprocess.run(
                ["git", "-C", str(repo_dir), "pull"],
                capture_output=True,
                text=True
            )
        
        print("📦 Installing DiffSynth-Studio...")
        subprocess.run(
            [sys.executable, "-m", "pip", "install", "-e", str(repo_dir)],
            capture_output=True,
            text=True,
            check=True
        )
        print("✅ DiffSynth-Studio installed successfully.")
        
        sys.path.insert(0, str(repo_dir))
        
        from diffsynth.pipelines.z_image import ZImagePipeline
        return True, "✅ DiffSynth-Studio installed successfully!"
        
    except subprocess.CalledProcessError as e:
        error_msg = f"❌ Installation failed: {e.stderr}"
        print(error_msg)
        return False, error_msg
    except Exception as e:
        error_msg = f"❌ Error during installation: {str(e)}"
        print(error_msg)
        return False, error_msg


# =============================================================================
# Pipeline Initialization
# =============================================================================

print("=" * 60)
print("  Z-Image-i2L Gradio Demo - Initializing")
print("=" * 60)
print()

# Step 1: Install DiffSynth-Studio
print("🔍 Step 1: Checking DiffSynth-Studio installation...")
success, message = install_diffsynth_studio()
print(message)

if not success:
    raise RuntimeError("Failed to install DiffSynth-Studio. Cannot continue.")

# Step 2: Download HuggingFace models
print()
print("🔍 Step 2: Downloading models from HuggingFace...")
print(f"   Models directory: {MODELS_DIR.absolute()}")
downloaded_paths = download_hf_models(MODELS_DIR)

# Import required modules
from diffsynth.pipelines.z_image import (
    ZImagePipeline, ModelConfig,
    ZImageUnit_Image2LoRAEncode, ZImageUnit_Image2LoRADecode
)
from safetensors.torch import save_file, load_file

# Step 3: Configure VRAM settings
print()
print("⚙️  Step 3: Configuring VRAM settings...")
vram_config = {
    "offload_dtype": torch.bfloat16,
    "offload_device": "cuda",
    "onload_dtype": torch.bfloat16,
    "onload_device": "cuda",
    "preparing_dtype": torch.bfloat16,
    "preparing_device": "cuda",
    "computation_dtype": torch.bfloat16,
    "computation_device": "cuda",
}

# Step 4: Resolve local model paths
print()
print("📂 Step 4: Resolving model paths...")

# Z-Image transformer
zimage_path = MODELS_DIR / "Tongyi-MAI" / "Z-Image"
zimage_transformer_files = get_model_files(zimage_path, "transformer/*.safetensors")

# Z-Image-Turbo
zimage_turbo_path = MODELS_DIR / "Tongyi-MAI" / "Z-Image-Turbo"
text_encoder_files = get_model_files(zimage_turbo_path, "text_encoder/*.safetensors")
vae_file = get_model_files(zimage_turbo_path, "vae/diffusion_pytorch_model.safetensors")
tokenizer_path = zimage_turbo_path / "tokenizer"

# General Image Encoders
encoders_path = MODELS_DIR / "DiffSynth-Studio" / "General-Image-Encoders"
siglip_file = get_model_files(encoders_path, "SigLIP2-G384/model.safetensors")
dino_file = get_model_files(encoders_path, "DINOv3-7B/model.safetensors")

# Z-Image-i2L from HuggingFace
zimage_i2l_path = MODELS_DIR / "DiffSynth-Studio" / "Z-Image-i2L"
zimage_i2l_file = get_model_files(zimage_i2l_path, "model.safetensors")

print(f"   Z-Image transformer: {len(zimage_transformer_files)} file(s)")
print(f"   Text encoder: {len(text_encoder_files)} file(s)")
print(f"   VAE: {len(vae_file)} file(s)")
print(f"   Tokenizer: {tokenizer_path}")
print(f"   SigLIP2: {len(siglip_file)} file(s)")
print(f"   DINOv3: {len(dino_file)} file(s)")
print(f"   Z-Image-i2L: {len(zimage_i2l_file)} file(s)")

# Validate files
missing = []
if not zimage_transformer_files: missing.append("Z-Image transformer")
if not text_encoder_files: missing.append("Text encoder")
if not vae_file: missing.append("VAE")
if not tokenizer_path.exists(): missing.append("Tokenizer")
if not siglip_file: missing.append("SigLIP2")
if not dino_file: missing.append("DINOv3")
if not zimage_i2l_file: missing.append("Z-Image-i2L")

if missing:
    raise FileNotFoundError(f"Missing model files: {', '.join(missing)}")

# Step 5: Load pipeline
print()
print("🚀 Step 5: Loading Z-Image pipeline...")
print("   All models loaded from HuggingFace local paths")

model_configs = [
    # All models from HuggingFace - use path= for local files
    ModelConfig(path=zimage_transformer_files, **vram_config),
    ModelConfig(path=text_encoder_files),
    ModelConfig(path=vae_file),
    ModelConfig(path=siglip_file),
    ModelConfig(path=dino_file),
    ModelConfig(path=zimage_i2l_file),
]

pipe = ZImagePipeline.from_pretrained(
    torch_dtype=torch.bfloat16,
    device="cuda",
    model_configs=model_configs,
    tokenizer_config=ModelConfig(path=str(tokenizer_path)),
)

print()
print("✅ Pipeline loaded successfully!")
print("=" * 60)
print()


# =============================================================================
# Gradio Functions
# =============================================================================

@spaces.GPU(duration=120)
def image_to_lora(images, progress=gr.Progress()):
    """Convert input images to a LoRA model."""
    if images is None or len(images) == 0:
        return None, "❌ Please upload at least one image!"
    
    try:
        progress(0.1, desc="Processing images...")
        
        pil_images = []
        for img in images:
            if isinstance(img, str):
                pil_images.append(Image.open(img).convert("RGB"))
            elif isinstance(img, tuple):
                pil_images.append(Image.open(img[0]).convert("RGB"))
            else:
                pil_images.append(Image.fromarray(img).convert("RGB"))
        
        progress(0.3, desc="Encoding images to LoRA...")
        
        with torch.no_grad():
            embs = ZImageUnit_Image2LoRAEncode().process(pipe, image2lora_images=pil_images)
            progress(0.7, desc="Decoding LoRA weights...")
            lora = ZImageUnit_Image2LoRADecode().process(pipe, **embs)["lora"]
        
        progress(0.9, desc="Saving LoRA file...")
        
        temp_dir = tempfile.mkdtemp()
        lora_path = os.path.join(temp_dir, "generated_lora.safetensors")
        save_file(lora, lora_path)
        
        progress(1.0, desc="Done!")
        
        return lora_path, f"✅ LoRA generated successfully from {len(pil_images)} image(s)!"
    
    except Exception as e:
        return None, f"❌ Error generating LoRA: {str(e)}"


@spaces.GPU(duration=60)
def generate_image(
    lora_file,
    prompt,
    negative_prompt,
    seed,
    cfg_scale,
    sigma_shift,
    num_steps,
    progress=gr.Progress()
):
    """Generate an image using the created LoRA."""
    if lora_file is None:
        return None, "❌ Please generate or upload a LoRA file first!"
    
    try:
        progress(0.1, desc="Loading LoRA...")
        
        lora = load_file(lora_file)
        # Move LoRA tensors to CUDA with correct dtype
        lora = {k: v.to(device="cuda", dtype=torch.bfloat16) for k, v in lora.items()}
        
        progress(0.3, desc="Generating image...")
        
        image = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            seed=int(seed),
            cfg_scale=cfg_scale,
            num_inference_steps=int(num_steps),
            positive_only_lora=lora,
            sigma_shift=sigma_shift
        )
        
        progress(1.0, desc="Done!")
        
        return image, "✅ Image generated successfully!"
    
    except Exception as e:
        return None, f"❌ Error generating image: {str(e)}"


def create_demo():
    """Create the Gradio interface."""
    
    with gr.Blocks(
        title="Z-Image-i2L Demo",
        theme=gr.themes.Soft(),
        css=".gradio-container { max-width: 1200px !important; margin: 0 auto}"
    ) as demo:
        gr.Markdown("""
        # 🎨 Z-Image-i2L: Image to LoRA Demo
        
        > 💡 **Tip**: For best results, use 4-6 images with a consistent artistic style.
        """)
        
        with gr.Tabs():
            with gr.TabItem("📸 Step 1: Image to LoRA"):
                with gr.Row():
                    with gr.Column(scale=1):
                        input_gallery = gr.Gallery(
                            label="Upload Style Images (1-6 images)",
                            file_types=["image"],
                            columns=3,
                            height=300,
                            interactive=True
                        )
                        
                        gr.Markdown("""
                        **Guidelines:**
                        - Upload 1-6 images with a consistent style
                        - Higher quality images produce better results
                        - Mix of subjects helps generalization
                        """)
                        
                        generate_lora_btn = gr.Button("🎯 Generate LoRA", variant="primary")
                    
                    with gr.Column(scale=1):
                        lora_output = gr.File(
                            label="Generated LoRA File",
                            file_types=[".safetensors"],
                            interactive=False
                        )
                        lora_status = gr.Textbox(
                            label="Status",
                            interactive=False,
                            lines=2
                        )
            
            with gr.TabItem("🖼️ Step 2: Generate Images"):
                with gr.Row():
                    with gr.Column(scale=1):
                        lora_input = gr.File(
                            label="LoRA File (from Step 1 or upload)",
                            file_types=[".safetensors"]
                        )
                        
                        prompt = gr.Textbox(
                            label="Prompt",
                            placeholder="Describe what you want to generate...",
                            value="a cat",
                            lines=2
                        )
                        
                        with gr.Accordion("Negative Prompt", open=False):
                            negative_prompt = gr.Textbox(
                                label="Negative Prompt",
                                value=NEGATIVE_PROMPT_CN,
                                lines=3
                            )
                            with gr.Row():
                                use_cn_neg = gr.Button("Use Chinese", size="sm")
                                use_en_neg = gr.Button("Use English", size="sm")
                        
                        with gr.Accordion("Advanced Settings", open=False):
                            seed = gr.Number(label="Seed", value=0, precision=0)
                            cfg_scale = gr.Slider(label="CFG Scale", minimum=1, maximum=10, value=4, step=0.5)
                            sigma_shift = gr.Slider(label="Sigma Shift", minimum=1, maximum=15, value=8, step=1)
                            num_steps = gr.Slider(label="Steps", minimum=20, maximum=100, value=50, step=5)
                        
                        generate_btn = gr.Button("✨ Generate Image", variant="primary")
                    
                    with gr.Column(scale=1):
                        output_image = gr.Image(label="Generated Image", type="pil", height=512)
                        gen_status = gr.Textbox(label="Status", interactive=False, lines=2)
        
        gr.Markdown("""
        ---
        **Resources:** [Z-Image-i2L (HuggingFace)](https://huggingface.co/DiffSynth-Studio/Z-Image-i2L) | 
        [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio) |
        **Settings:** CFG=4, Sigma Shift=8, Steps=50
        """)
        
        # Event handlers
        generate_lora_btn.click(
            fn=image_to_lora,
            inputs=[input_gallery],
            outputs=[lora_output, lora_status]
        )
        
        lora_output.change(fn=lambda x: x, inputs=[lora_output], outputs=[lora_input])
        
        generate_btn.click(
            fn=generate_image,
            inputs=[lora_input, prompt, negative_prompt, seed, cfg_scale, sigma_shift, num_steps],
            outputs=[output_image, gen_status]
        )
        
        use_cn_neg.click(fn=lambda: NEGATIVE_PROMPT_CN, outputs=[negative_prompt])
        use_en_neg.click(fn=lambda: NEGATIVE_PROMPT_EN, outputs=[negative_prompt])
    
    return demo


if __name__ == "__main__":
    print("Starting Gradio server...")
    demo = create_demo()
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)