Update app.py

app.py

@@ -6,8 +6,9 @@ A web interface for the Z-Image-i2L model that converts images to LoRA models.
 Setup Instructions:
 1. Install dependencies:
    pip install -r requirements.txt
+   pip install huggingface_hub
 
-2. Run this demo (DiffSynth-Studio will be auto-installed):
+2. Run this demo (models will be auto-downloaded):
    python app.py
 
 Note: This requires a GPU with sufficient VRAM (recommended 24GB+)
@@ -22,29 +23,110 @@ 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/General-Image-Encoders
+    - Tongyi-MAI/Z-Image-Turbo  
+    - Tongyi-MAI/Z-Image
+    
+    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"],
+        },
+    ]
+    
+    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."""
-    # Check if already installed
     try:
         from diffsynth.pipelines.z_image import ZImagePipeline
         return True, "✅ DiffSynth-Studio is already installed."
     except ImportError:
         pass
     
-    # Define paths
     repo_dir = Path(__file__).parent / "DiffSynth-Studio"
     
     try:
-        # Clone the repository if it doesn't exist
         if not repo_dir.exists():
             print("📥 Cloning DiffSynth-Studio repository...")
-            result = subprocess.run(
+            subprocess.run(
                 ["git", "clone", "https://github.com/modelscope/DiffSynth-Studio.git", str(repo_dir)],
                 capture_output=True,
                 text=True,
@@ -53,15 +135,14 @@ def install_diffsynth_studio():
             print("✅ Repository cloned successfully.")
         else:
             print("📁 DiffSynth-Studio directory already exists, pulling latest...")
-            result = subprocess.run(
+            subprocess.run(
                 ["git", "-C", str(repo_dir), "pull"],
                 capture_output=True,
                 text=True
             )
         
-        # Install in editable mode
         print("📦 Installing DiffSynth-Studio...")
-        result = subprocess.run(
+        subprocess.run(
             [sys.executable, "-m", "pip", "install", "-e", str(repo_dir)],
             capture_output=True,
             text=True,
@@ -69,7 +150,6 @@ def install_diffsynth_studio():
         )
         print("✅ DiffSynth-Studio installed successfully.")
         
-        # Add to path and try importing again
         sys.path.insert(0, str(repo_dir))
         
         from diffsynth.pipelines.z_image import ZImagePipeline
@@ -86,22 +166,28 @@ def install_diffsynth_studio():
 
 
 # =============================================================================
-# Pipeline Initialization (runs at module load time)
+# Pipeline Initialization
 # =============================================================================
 
-print("=" * 50)
+print("=" * 60)
 print("  Z-Image-i2L Gradio Demo - Initializing")
-print("=" * 50)
+print("=" * 60)
 print()
 
-# Ensure DiffSynth-Studio is installed
-print("🔍 Checking DiffSynth-Studio installation...")
+# 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,
@@ -109,8 +195,9 @@ from diffsynth.pipelines.z_image import (
 )
 from safetensors.torch import save_file, load_file
 
-# Configure VRAM settings
-print("⚙️  Configuring VRAM settings...")
+# Step 3: Configure VRAM settings
+print()
+print("⚙️  Step 3: Configuring VRAM settings...")
 vram_config = {
     "offload_dtype": torch.bfloat16,
     "offload_device": "cuda",
@@ -122,27 +209,75 @@ vram_config = {
     "computation_device": "cuda",
 }
 
-# Load the pipeline
-print("🚀 Loading Z-Image pipeline...")
-print("   This may take a few minutes on first run (downloading models)...")
+# 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")
+
+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: ModelScope (auto-download)")
+
+# 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 missing:
+    raise FileNotFoundError(f"Missing model files: {', '.join(missing)}")
+
+# Step 5: Load pipeline
+print()
+print("🚀 Step 5: Loading Z-Image pipeline...")
+print("   HuggingFace models: loaded from local paths")
+print("   Z-Image-i2L: loading from ModelScope...")
+
+model_configs = [
+    # HuggingFace models - 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),
+    # ModelScope only - use model_id= for remote download
+    ModelConfig(model_id="DiffSynth-Studio/Z-Image-i2L", origin_file_pattern="model.safetensors"),
+]
 
 pipe = ZImagePipeline.from_pretrained(
     torch_dtype=torch.bfloat16,
     device="cuda",
-    model_configs=[
-        ModelConfig(model_id="Tongyi-MAI/Z-Image", origin_file_pattern="transformer/*.safetensors", **vram_config),
-        ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="text_encoder/*.safetensors"),
-        ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="vae/diffusion_pytorch_model.safetensors"),
-        ModelConfig(model_id="DiffSynth-Studio/General-Image-Encoders", origin_file_pattern="SigLIP2-G384/model.safetensors"),
-        ModelConfig(model_id="DiffSynth-Studio/General-Image-Encoders", origin_file_pattern="DINOv3-7B/model.safetensors"),
-        ModelConfig(model_id="DiffSynth-Studio/Z-Image-i2L", origin_file_pattern="model.safetensors"),
-    ],
-    tokenizer_config=ModelConfig(model_id="Tongyi-MAI/Z-Image-Turbo", origin_file_pattern="tokenizer/"),
+    model_configs=model_configs,
+    tokenizer_config=ModelConfig(path=str(tokenizer_path)),
 )
 
+print()
 print("✅ Pipeline loaded successfully!")
+print("=" * 60)
 print()
 
+
 # =============================================================================
 # Gradio Functions
 # =============================================================================
@@ -156,13 +291,11 @@ def image_to_lora(images, progress=gr.Progress()):
     try:
         progress(0.1, desc="Processing images...")
         
-        # Convert uploaded images to PIL
         pil_images = []
         for img in images:
             if isinstance(img, str):
                 pil_images.append(Image.open(img).convert("RGB"))
             elif isinstance(img, tuple):
-                # Gradio gallery returns tuples (filepath, caption)
                 pil_images.append(Image.open(img[0]).convert("RGB"))
             else:
                 pil_images.append(Image.fromarray(img).convert("RGB"))
@@ -176,14 +309,13 @@ def image_to_lora(images, progress=gr.Progress()):
         
         progress(0.9, desc="Saving LoRA file...")
         
-        # Save to temporary 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)} images!"
+        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)}"
@@ -235,10 +367,7 @@ def create_demo():
     with gr.Blocks(
         title="Z-Image-i2L Demo",
         theme=gr.themes.Soft(),
-        css="""
-        .gradio-container { max-width: 1200px !important; }
-        .status-box { padding: 10px; border-radius: 5px; margin: 10px 0; }
-        """
+        css=".gradio-container { max-width: 1200px !important; }"
     ) as demo:
         gr.Markdown("""
         # 🎨 Z-Image-i2L: Image to LoRA Demo
@@ -251,12 +380,9 @@ def create_demo():
         3. **Generate Images**: Use the LoRA to create new images with your style
         
         > 💡 **Tip**: For best results, use 4-6 images with a consistent artistic style.
-        
-        ✅ **Pipeline is pre-loaded and ready to use!**
         """)
         
         with gr.Tabs():
-            # Tab 1: Image to LoRA
             with gr.TabItem("📸 Step 1: Image to LoRA"):
                 with gr.Row():
                     with gr.Column(scale=1):
@@ -269,10 +395,10 @@ def create_demo():
                         )
                         
                         gr.Markdown("""
-                        **Guidelines for input images:**
+                        **Guidelines:**
                         - Upload 1-6 images with a consistent style
                         - Higher quality images produce better results
-                        - Mix of subjects (people, objects, scenes) helps generalization
+                        - Mix of subjects helps generalization
                         """)
                         
                         generate_lora_btn = gr.Button("🎯 Generate LoRA", variant="primary")
@@ -284,12 +410,11 @@ def create_demo():
                             interactive=False
                         )
                         lora_status = gr.Textbox(
-                            label="LoRA Generation Status",
+                            label="Status",
                             interactive=False,
                             lines=2
                         )
             
-            # Tab 2: Generate Images
             with gr.TabItem("🖼️ Step 2: Generate Images"):
                 with gr.Row():
                     with gr.Column(scale=1):
@@ -316,59 +441,22 @@ def create_demo():
                                 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="Number of Steps",
-                                minimum=20,
-                                maximum=100,
-                                value=50,
-                                step=5
-                            )
+                            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="Generation Status",
-                            interactive=False,
-                            lines=2
-                        )
+                        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 Model](https://modelscope.cn/models/DiffSynth-Studio/Z-Image-i2L)
-        - [DiffSynth-Studio GitHub](https://github.com/modelscope/DiffSynth-Studio)
-        - [Online Demo](https://modelscope.cn/studios/DiffSynth-Studio/Z-Image-i2L)
-        
-        ### ⚙️ Recommended Settings
-        - **CFG Scale**: 4
-        - **Sigma Shift**: 8
-        - **Steps**: 50
-        - Use negative prompts for better quality
+        **Resources:** [Z-Image-i2L](https://modelscope.cn/models/DiffSynth-Studio/Z-Image-i2L) | 
+        [DiffSynth-Studio](https://github.com/modelscope/DiffSynth-Studio) |
+        **Settings:** CFG=4, Sigma Shift=8, Steps=50
         """)
         
         # Event handlers
@@ -378,47 +466,21 @@ def create_demo():
             outputs=[lora_output, lora_status]
         )
         
-        # Auto-update lora_input when lora is generated
-        lora_output.change(
-            fn=lambda x: x,
-            inputs=[lora_output],
-            outputs=[lora_input]
-        )
+        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
-            ],
+            inputs=[lora_input, prompt, negative_prompt, seed, cfg_scale, sigma_shift, num_steps],
             outputs=[output_image, gen_status]
         )
         
-        # Negative prompt presets
-        use_cn_neg.click(
-            fn=lambda: NEGATIVE_PROMPT_CN,
-            outputs=[negative_prompt]
-        )
-        use_en_neg.click(
-            fn=lambda: NEGATIVE_PROMPT_EN,
-            outputs=[negative_prompt]
-        )
+        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...")
-    print()
-    
     demo = create_demo()
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False
-    )
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)