Enhanced: Real-ESRGAN upscaler, GFPGAN face restoration, multi-stage detailer

.gitattributes

@@ -1,38 +1,38 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
-disaster_girl.jpg filter=lfs diff=lfs merge=lfs -text
-grumpy.png filter=lfs diff=lfs merge=lfs -text
-wednesday.png filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+disaster_girl.jpg filter=lfs diff=lfs merge=lfs -text
+grumpy.png filter=lfs diff=lfs merge=lfs -text
+wednesday.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,58 +1,58 @@
-# IDE configuration files
-pyrightconfig.json
-.vscode/
-
-# Generated outputs
-outputs/
-
-# Custom folder (user-specific files)
-custom/
-
-# Python
-__pycache__/
-*.py[cod]
-*$py.class
-*.so
-.Python
-env/
-venv/
-ENV/
-build/
-develop-eggs/
-dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-*.egg-info/
-.installed.cfg
-*.egg
-
-# Model files (if you don't want to track large files)
-*.pt
-*.pth
-*.bin
-*.safetensors
-
-# OS files
-.DS_Store
-Thumbs.db
-
-# Testing
-.pytest_cache/
-.coverage
-htmlcov/
-
-# Logs
-*.log
-
-# Backup files
-*.from_server
-*.backup
-*.bak
-.claude/
+# IDE configuration files
+pyrightconfig.json
+.vscode/
+
+# Generated outputs
+outputs/
+
+# Custom folder (user-specific files)
+custom/
+
+# Python
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+env/
+venv/
+ENV/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# Model files (if you don't want to track large files)
+*.pt
+*.pth
+*.bin
+*.safetensors
+
+# OS files
+.DS_Store
+Thumbs.db
+
+# Testing
+.pytest_cache/
+.coverage
+htmlcov/
+
+# Logs
+*.log
+
+# Backup files
+*.from_server
+*.backup
+*.bak
+.claude/

AGENT.md

@@ -1,169 +1,169 @@
-# AGENT.md
-
-This file provides guidance to AI coding agents when working with code in this repository.
-
-## Project Overview
-
-This is a Gradio Space that implements "Next Scene" cinematic image generation using Qwen-Image-Edit-2509 with LoRA fine-tuning. The application generates visually progressive image sequences with natural cinematic transitions from frame to frame, optimized for fast 4-step inference.
-
-**Key Model Components:**
-- Base model: `Qwen/Qwen-Image-Edit-2509` (image editing diffusion model)
-- Accelerated transformer: `Phr00t/Qwen-Image-Edit-Rapid-AIO` v23 single-file checkpoint (4-step optimized variant)
-- LoRA adapter: `lovis93/next-scene-qwen-image-lora-2509` (cinematic progression fine-tune)
-- Text encoder: `Qwen2.5-VL-72B-Instruct` (via Hugging Face InferenceClient for prompt enhancement)
-
-## Running the Application
-
-**Start the Gradio interface:**
-```bash
-python app.py
-```
-
-**Install dependencies:**
-```bash
-pip install -r requirements.txt
-```
-
-The app requires GPU access. It uses the `@spaces.GPU` decorator for Hugging Face Spaces zero-GPU allocation.
-
-## Architecture
-
-### Pipeline Flow
-
-1. **Input Processing** (`app.py:infer`):
-   - Accepts input images via Gradio Gallery (filepath-based)
-   - Optional prompt rewriting using `Qwen2.5-VL-72B-Instruct` API
-   - Automatic "Next Scene" prompt generation from images
-
-2. **Image Generation** (`qwenimage/pipeline_qwenimage_edit_plus.py`):
-   - Custom pipeline extending `DiffusionPipeline`
-   - Encodes images using VAE at 1024x1024 for latents
-   - Encodes conditioning images at 384x384 for text encoder
-   - Packs latents into 2x2 patches (latent dims must be divisible by 2)
-   - Uses `FlowMatchEulerDiscreteScheduler` for denoising
-
-3. **Optimization** (`optimization.py`):
-   - Ahead-of-time (AOT) compilation using `torch.export` and `spaces.aoti_compile`
-   - Dynamic shapes for variable sequence lengths
-   - Custom inductor configs for performance (max_autotune, cudagraphs)
-   - FlashAttention 3 integration via `QwenDoubleStreamAttnProcessorFA3`
-
-4. **Output Handling**:
-   - Saves outputs to `outputs/` directory with unique timestamps
-   - Maintains 20-image history gallery
-   - Optional video generation through an owned Gradio Space configured by `VIDEO_SPACE_ID`
-
-### Custom QwenImage Components
-
-**Location:** `qwenimage/` package
-
-- `pipeline_qwenimage_edit_plus.py` - Main diffusion pipeline with LoRA support
-- `transformer_qwenimage.py` - Custom transformer model with cache management
-- `qwen_fa3_processor.py` - FlashAttention 3 attention processor
-
-**Key architectural features:**
-- Latent packing/unpacking for 2x2 patch processing
-- Multi-image conditioning support
-- True CFG (classifier-free guidance) with separate pos/neg paths
-- Dual-stream attention with rotary embeddings
-- Cache contexts for conditional/unconditional forward passes
-
-### Prompt Engineering
-
-**Two-stage prompt system:**
-
-1. **Edit Instruction Rewriter** (`SYSTEM_PROMPT`):
-   - Normalizes user prompts into professional editing instructions
-   - Handles text replacement (requires quotes), object manipulation, style transfer
-   - Used when `rewrite_prompt=True` checkbox is enabled
-
-2. **Next Scene Generator** (`NEXT_SCENE_SYSTEM_PROMPT`):
-   - Automatically suggests cinematic camera movements
-   - Focus on visual progression (dolly, pan, zoom, framing changes)
-   - Auto-triggers when input images change
-
-Both use `Qwen2.5-VL-72B-Instruct` via Hugging Face InferenceClient with Nebius provider. Requires `HF_TOKEN` environment variable.
-
-## Important Implementation Details
-
-### Image Dimension Handling
-
-Images are automatically resized based on `calculate_dimensions()` function:
-- VAE images: resized to maintain 1024×1024 area (1,048,576 pixels)
-- Condition images: resized to maintain 384×384 area (147,456 pixels)
-- Output dimensions must be divisible by 16 (vae_scale_factor × 2)
-- Height/width default to `None` which auto-calculates from input aspect ratio
-
-### LoRA Integration
-
-The pipeline fuses the "next-scene" LoRA adapter at initialization:
-```python
-pipe.load_lora_weights("lovis93/next-scene-qwen-image-lora-2509", ...)
-pipe.set_adapters(["next-scene"], adapter_weights=[1.])
-pipe.fuse_lora(adapter_names=["next-scene"], lora_scale=1.)
-pipe.unload_lora_weights()
-```
-
-After fusion, the adapter weights are merged into the base model and cannot be unfused.
-
-### Video Generation Integration
-
-The `turn_into_video()` function:
-- Connects only when `VIDEO_SPACE_ID` is configured
-- Requires first input image and last output image
-- Uses the original prompt (or "smooth cinematic transition" fallback)
-- Returns video path for display
-
-### Gradio Gallery Format
-
-Input/output galleries use `type="filepath"` (string paths) rather than PIL Image tuples. Helper functions handle format compatibility for legacy tuple support.
-
-## Environment Variables
-
-- `HF_TOKEN` - Required for Qwen2.5-VL API access (prompt rewriting/generation)
-
-## File Outputs
-
-Generated images are saved to `outputs/` directory with format:
-```
-output_{seed}_{index}_{timestamp_ms}.png
-```
-
-## Local Development and API Testing
-
-The `custom/` directory is fully gitignored and used for local development files. Specifically, it contains:
-
-- **API client scripts** - For testing the Gradio Space remotely via API after deployment to Hugging Face
-- **`API_GUIDE.txt`** - Auto-generated Gradio API documentation showing endpoint signatures and example usage
-- **Local testing environments** - Virtual environments or test data that shouldn't be committed
-
-**API Integration Pattern:**
-Once the Space is deployed to Hugging Face, you can interact with it programmatically using `gradio_client`:
-
-```python
-from gradio_client import Client, handle_file
-
-client = Client("clopeux/Pro-Realism-Edit-Studio")
-result = client.predict(
-    images=[],
-    prompt="Next Scene: Camera dollies forward...",
-    seed=42,
-    randomize_seed=False,
-    true_guidance_scale=1.0,
-    num_inference_steps=4,
-    height=1024,
-    width=1024,
-    rewrite_prompt=False,
-    api_name="/infer"
-)
-```
-
-The `custom/API_GUIDE.txt` contains full documentation of all available endpoints including `/infer`, `/turn_into_video`, `/suggest_next_scene_prompt`, and utility functions.
-
-## Development Notes
-
-- The model loads on startup and applies AOT compilation during first inference
-- Compilation uses dynamic shapes to support variable text/image sequence lengths
-- The transformer uses custom cache contexts ("cond"/"uncond") to optimize CFG passes
-- True CFG applies norm-based rescaling: `comb_pred * (cond_norm / noise_norm)`
-- FlashAttention 3 processor must be set before compilation
+# AGENT.md
+
+This file provides guidance to AI coding agents when working with code in this repository.
+
+## Project Overview
+
+This is a Gradio Space that implements "Next Scene" cinematic image generation using Qwen-Image-Edit-2509 with LoRA fine-tuning. The application generates visually progressive image sequences with natural cinematic transitions from frame to frame, optimized for fast 4-step inference.
+
+**Key Model Components:**
+- Base model: `Qwen/Qwen-Image-Edit-2509` (image editing diffusion model)
+- Accelerated transformer: `Phr00t/Qwen-Image-Edit-Rapid-AIO` v23 single-file checkpoint (4-step optimized variant)
+- LoRA adapter: `lovis93/next-scene-qwen-image-lora-2509` (cinematic progression fine-tune)
+- Text encoder: `Qwen2.5-VL-72B-Instruct` (via Hugging Face InferenceClient for prompt enhancement)
+
+## Running the Application
+
+**Start the Gradio interface:**
+```bash
+python app.py
+```
+
+**Install dependencies:**
+```bash
+pip install -r requirements.txt
+```
+
+The app requires GPU access. It uses the `@spaces.GPU` decorator for Hugging Face Spaces zero-GPU allocation.
+
+## Architecture
+
+### Pipeline Flow
+
+1. **Input Processing** (`app.py:infer`):
+   - Accepts input images via Gradio Gallery (filepath-based)
+   - Optional prompt rewriting using `Qwen2.5-VL-72B-Instruct` API
+   - Automatic "Next Scene" prompt generation from images
+
+2. **Image Generation** (`qwenimage/pipeline_qwenimage_edit_plus.py`):
+   - Custom pipeline extending `DiffusionPipeline`
+   - Encodes images using VAE at 1024x1024 for latents
+   - Encodes conditioning images at 384x384 for text encoder
+   - Packs latents into 2x2 patches (latent dims must be divisible by 2)
+   - Uses `FlowMatchEulerDiscreteScheduler` for denoising
+
+3. **Optimization** (`optimization.py`):
+   - Ahead-of-time (AOT) compilation using `torch.export` and `spaces.aoti_compile`
+   - Dynamic shapes for variable sequence lengths
+   - Custom inductor configs for performance (max_autotune, cudagraphs)
+   - FlashAttention 3 integration via `QwenDoubleStreamAttnProcessorFA3`
+
+4. **Output Handling**:
+   - Saves outputs to `outputs/` directory with unique timestamps
+   - Maintains 20-image history gallery
+   - Optional video generation through an owned Gradio Space configured by `VIDEO_SPACE_ID`
+
+### Custom QwenImage Components
+
+**Location:** `qwenimage/` package
+
+- `pipeline_qwenimage_edit_plus.py` - Main diffusion pipeline with LoRA support
+- `transformer_qwenimage.py` - Custom transformer model with cache management
+- `qwen_fa3_processor.py` - FlashAttention 3 attention processor
+
+**Key architectural features:**
+- Latent packing/unpacking for 2x2 patch processing
+- Multi-image conditioning support
+- True CFG (classifier-free guidance) with separate pos/neg paths
+- Dual-stream attention with rotary embeddings
+- Cache contexts for conditional/unconditional forward passes
+
+### Prompt Engineering
+
+**Two-stage prompt system:**
+
+1. **Edit Instruction Rewriter** (`SYSTEM_PROMPT`):
+   - Normalizes user prompts into professional editing instructions
+   - Handles text replacement (requires quotes), object manipulation, style transfer
+   - Used when `rewrite_prompt=True` checkbox is enabled
+
+2. **Next Scene Generator** (`NEXT_SCENE_SYSTEM_PROMPT`):
+   - Automatically suggests cinematic camera movements
+   - Focus on visual progression (dolly, pan, zoom, framing changes)
+   - Auto-triggers when input images change
+
+Both use `Qwen2.5-VL-72B-Instruct` via Hugging Face InferenceClient with Nebius provider. Requires `HF_TOKEN` environment variable.
+
+## Important Implementation Details
+
+### Image Dimension Handling
+
+Images are automatically resized based on `calculate_dimensions()` function:
+- VAE images: resized to maintain 1024×1024 area (1,048,576 pixels)
+- Condition images: resized to maintain 384×384 area (147,456 pixels)
+- Output dimensions must be divisible by 16 (vae_scale_factor × 2)
+- Height/width default to `None` which auto-calculates from input aspect ratio
+
+### LoRA Integration
+
+The pipeline fuses the "next-scene" LoRA adapter at initialization:
+```python
+pipe.load_lora_weights("lovis93/next-scene-qwen-image-lora-2509", ...)
+pipe.set_adapters(["next-scene"], adapter_weights=[1.])
+pipe.fuse_lora(adapter_names=["next-scene"], lora_scale=1.)
+pipe.unload_lora_weights()
+```
+
+After fusion, the adapter weights are merged into the base model and cannot be unfused.
+
+### Video Generation Integration
+
+The `turn_into_video()` function:
+- Connects only when `VIDEO_SPACE_ID` is configured
+- Requires first input image and last output image
+- Uses the original prompt (or "smooth cinematic transition" fallback)
+- Returns video path for display
+
+### Gradio Gallery Format
+
+Input/output galleries use `type="filepath"` (string paths) rather than PIL Image tuples. Helper functions handle format compatibility for legacy tuple support.
+
+## Environment Variables
+
+- `HF_TOKEN` - Required for Qwen2.5-VL API access (prompt rewriting/generation)
+
+## File Outputs
+
+Generated images are saved to `outputs/` directory with format:
+```
+output_{seed}_{index}_{timestamp_ms}.png
+```
+
+## Local Development and API Testing
+
+The `custom/` directory is fully gitignored and used for local development files. Specifically, it contains:
+
+- **API client scripts** - For testing the Gradio Space remotely via API after deployment to Hugging Face
+- **`API_GUIDE.txt`** - Auto-generated Gradio API documentation showing endpoint signatures and example usage
+- **Local testing environments** - Virtual environments or test data that shouldn't be committed
+
+**API Integration Pattern:**
+Once the Space is deployed to Hugging Face, you can interact with it programmatically using `gradio_client`:
+
+```python
+from gradio_client import Client, handle_file
+
+client = Client("clopeux/Pro-Realism-Edit-Studio")
+result = client.predict(
+    images=[],
+    prompt="Next Scene: Camera dollies forward...",
+    seed=42,
+    randomize_seed=False,
+    true_guidance_scale=1.0,
+    num_inference_steps=4,
+    height=1024,
+    width=1024,
+    rewrite_prompt=False,
+    api_name="/infer"
+)
+```
+
+The `custom/API_GUIDE.txt` contains full documentation of all available endpoints including `/infer`, `/turn_into_video`, `/suggest_next_scene_prompt`, and utility functions.
+
+## Development Notes
+
+- The model loads on startup and applies AOT compilation during first inference
+- Compilation uses dynamic shapes to support variable text/image sequence lengths
+- The transformer uses custom cache contexts ("cond"/"uncond") to optimize CFG passes
+- True CFG applies norm-based rescaling: `comb_pred * (cond_norm / noise_norm)`
+- FlashAttention 3 processor must be set before compilation

CLAUDE.md

@@ -1,154 +1,154 @@
-# CLAUDE.md
-
-This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
-
-## Project Overview
-
-This is a Gradio Space that implements "Next Scene" cinematic image generation using Qwen-Image-Edit-2511 with LoRA fine-tuning. The application generates visually progressive image sequences with natural cinematic transitions from frame to frame, optimized for fast 4-step inference.
-
-**Key Model Components:**
-- Base model: `Qwen/Qwen-Image-Edit-2511` (image editing diffusion model)
-- Accelerated transformer: `Phr00t/Qwen-Image-Edit-Rapid-AIO` v23 single-file checkpoint (4-step optimized)
-- LoRA adapter: `lovis93/next-scene-qwen-image-lora-2509` (cinematic progression fine-tune, trained on 2509)
-
-## Running the Application
-
-**Start the Gradio interface:**
-```bash
-python app.py
-```
-
-**Install dependencies:**
-```bash
-pip install -r requirements.txt
-```
-
-The app requires GPU access. It uses the `@spaces.GPU` decorator for Hugging Face Spaces zero-GPU allocation.
-
-## Architecture
-
-### Pipeline Flow
-
-1. **Input Processing** (`app.py:infer`):
-   - Accepts input images via Gradio Gallery (filepath-based)
-   - Uses user-provided prompts directly without modification
-
-2. **Image Generation** (`qwenimage/pipeline_qwenimage_edit_plus.py`):
-   - Custom pipeline extending `DiffusionPipeline`
-   - Encodes images using VAE at 1024x1024 for latents
-   - Encodes conditioning images at 384x384 for text encoder
-   - Packs latents into 2x2 patches (latent dims must be divisible by 2)
-   - Uses `FlowMatchEulerDiscreteScheduler` for denoising
-
-3. **Optimization** (`optimization.py`):
-   - Ahead-of-time (AOT) compilation using `torch.export` and `spaces.aoti_compile`
-   - Dynamic shapes for variable sequence lengths
-   - Custom inductor configs for performance (max_autotune, cudagraphs)
-   - FlashAttention 3 integration via `QwenDoubleStreamAttnProcessorFA3`
-
-4. **Output Handling**:
-   - Saves outputs to `outputs/` directory with unique timestamps
-   - Maintains 20-image history gallery
-   - Optional video generation through an owned Gradio Space configured by `VIDEO_SPACE_ID`
-
-### Custom QwenImage Components
-
-**Location:** `qwenimage/` package
-
-- `pipeline_qwenimage_edit_plus.py` - Main diffusion pipeline with LoRA support
-- `transformer_qwenimage.py` - Custom transformer model with cache management
-- `qwen_fa3_processor.py` - FlashAttention 3 attention processor
-
-**Key architectural features:**
-- Latent packing/unpacking for 2x2 patch processing
-- Multi-image conditioning support
-- True CFG (classifier-free guidance) with separate pos/neg paths
-- Dual-stream attention with rotary embeddings
-- Cache contexts for conditional/unconditional forward passes
-
-### Prompt Handling
-
-The application uses user-provided prompts directly without any preprocessing, rewriting, or AI-based enhancement. Users have full control over the exact prompt text that gets passed to the diffusion model.
-
-## Important Implementation Details
-
-### Image Dimension Handling
-
-Images are automatically resized based on `calculate_dimensions()` function:
-- VAE images: resized to maintain 1024×1024 area (1,048,576 pixels)
-- Condition images: resized to maintain 384×384 area (147,456 pixels)
-- Output dimensions must be divisible by 16 (vae_scale_factor × 2)
-- Height/width default to `None` which auto-calculates from input aspect ratio
-
-### LoRA Integration
-
-The pipeline fuses the "next-scene" LoRA adapter at initialization:
-```python
-pipe.load_lora_weights("lovis93/next-scene-qwen-image-lora-2509", ...)
-pipe.set_adapters(["next-scene"], adapter_weights=[1.])
-pipe.fuse_lora(adapter_names=["next-scene"], lora_scale=1.)
-pipe.unload_lora_weights()
-```
-
-After fusion, the adapter weights are merged into the base model and cannot be unfused.
-
-### Video Generation Integration
-
-The `turn_into_video()` function:
-- Connects only when `VIDEO_SPACE_ID` is configured
-- Requires first input image and last output image
-- Uses the original prompt (or "smooth cinematic transition" fallback)
-- Returns video path for display
-
-### Gradio Gallery Format
-
-Input/output galleries use `type="filepath"` (string paths) rather than PIL Image tuples. Helper functions handle format compatibility for legacy tuple support.
-
-## Environment Variables
-
-No environment variables are required for basic operation. The application runs entirely with local models.
-
-## File Outputs
-
-Generated images are saved to `outputs/` directory with format:
-```
-output_{seed}_{index}_{timestamp_ms}.png
-```
-
-## Local Development and API Testing
-
-The `custom/` directory is fully gitignored and used for local development files. Specifically, it contains:
-
-- **API client scripts** - For testing the Gradio Space remotely via API after deployment to Hugging Face
-- **`API_GUIDE.txt`** - Auto-generated Gradio API documentation showing endpoint signatures and example usage
-- **Local testing environments** - Virtual environments or test data that shouldn't be committed
-
-**API Integration Pattern:**
-Once the Space is deployed to Hugging Face, you can interact with it programmatically using `gradio_client`:
-
-```python
-from gradio_client import Client, handle_file
-
-client = Client("clopeux/Pro-Realism-Edit-Studio")
-result = client.predict(
-    images=[],
-    prompt="Camera dollies forward, revealing more of the scene",
-    seed=42,
-    randomize_seed=False,
-    true_guidance_scale=1.0,
-    num_inference_steps=4,
-    height=1024,
-    width=1024,
-    api_name="/infer"
-)
-```
-
-The `custom/API_GUIDE.txt` contains full documentation of all available endpoints including `/infer`, `/turn_into_video`, and utility functions.
-
-## Development Notes
-
-- The model loads on startup and applies AOT compilation during first inference
-- Compilation uses dynamic shapes to support variable text/image sequence lengths
-- The transformer uses custom cache contexts ("cond"/"uncond") to optimize CFG passes
-- True CFG applies norm-based rescaling: `comb_pred * (cond_norm / noise_norm)`
-- FlashAttention 3 processor must be set before compilation
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Project Overview
+
+This is a Gradio Space that implements "Next Scene" cinematic image generation using Qwen-Image-Edit-2511 with LoRA fine-tuning. The application generates visually progressive image sequences with natural cinematic transitions from frame to frame, optimized for fast 4-step inference.
+
+**Key Model Components:**
+- Base model: `Qwen/Qwen-Image-Edit-2511` (image editing diffusion model)
+- Accelerated transformer: `Phr00t/Qwen-Image-Edit-Rapid-AIO` v23 single-file checkpoint (4-step optimized)
+- LoRA adapter: `lovis93/next-scene-qwen-image-lora-2509` (cinematic progression fine-tune, trained on 2509)
+
+## Running the Application
+
+**Start the Gradio interface:**
+```bash
+python app.py
+```
+
+**Install dependencies:**
+```bash
+pip install -r requirements.txt
+```
+
+The app requires GPU access. It uses the `@spaces.GPU` decorator for Hugging Face Spaces zero-GPU allocation.
+
+## Architecture
+
+### Pipeline Flow
+
+1. **Input Processing** (`app.py:infer`):
+   - Accepts input images via Gradio Gallery (filepath-based)
+   - Uses user-provided prompts directly without modification
+
+2. **Image Generation** (`qwenimage/pipeline_qwenimage_edit_plus.py`):
+   - Custom pipeline extending `DiffusionPipeline`
+   - Encodes images using VAE at 1024x1024 for latents
+   - Encodes conditioning images at 384x384 for text encoder
+   - Packs latents into 2x2 patches (latent dims must be divisible by 2)
+   - Uses `FlowMatchEulerDiscreteScheduler` for denoising
+
+3. **Optimization** (`optimization.py`):
+   - Ahead-of-time (AOT) compilation using `torch.export` and `spaces.aoti_compile`
+   - Dynamic shapes for variable sequence lengths
+   - Custom inductor configs for performance (max_autotune, cudagraphs)
+   - FlashAttention 3 integration via `QwenDoubleStreamAttnProcessorFA3`
+
+4. **Output Handling**:
+   - Saves outputs to `outputs/` directory with unique timestamps
+   - Maintains 20-image history gallery
+   - Optional video generation through an owned Gradio Space configured by `VIDEO_SPACE_ID`
+
+### Custom QwenImage Components
+
+**Location:** `qwenimage/` package
+
+- `pipeline_qwenimage_edit_plus.py` - Main diffusion pipeline with LoRA support
+- `transformer_qwenimage.py` - Custom transformer model with cache management
+- `qwen_fa3_processor.py` - FlashAttention 3 attention processor
+
+**Key architectural features:**
+- Latent packing/unpacking for 2x2 patch processing
+- Multi-image conditioning support
+- True CFG (classifier-free guidance) with separate pos/neg paths
+- Dual-stream attention with rotary embeddings
+- Cache contexts for conditional/unconditional forward passes
+
+### Prompt Handling
+
+The application uses user-provided prompts directly without any preprocessing, rewriting, or AI-based enhancement. Users have full control over the exact prompt text that gets passed to the diffusion model.
+
+## Important Implementation Details
+
+### Image Dimension Handling
+
+Images are automatically resized based on `calculate_dimensions()` function:
+- VAE images: resized to maintain 1024×1024 area (1,048,576 pixels)
+- Condition images: resized to maintain 384×384 area (147,456 pixels)
+- Output dimensions must be divisible by 16 (vae_scale_factor × 2)
+- Height/width default to `None` which auto-calculates from input aspect ratio
+
+### LoRA Integration
+
+The pipeline fuses the "next-scene" LoRA adapter at initialization:
+```python
+pipe.load_lora_weights("lovis93/next-scene-qwen-image-lora-2509", ...)
+pipe.set_adapters(["next-scene"], adapter_weights=[1.])
+pipe.fuse_lora(adapter_names=["next-scene"], lora_scale=1.)
+pipe.unload_lora_weights()
+```
+
+After fusion, the adapter weights are merged into the base model and cannot be unfused.
+
+### Video Generation Integration
+
+The `turn_into_video()` function:
+- Connects only when `VIDEO_SPACE_ID` is configured
+- Requires first input image and last output image
+- Uses the original prompt (or "smooth cinematic transition" fallback)
+- Returns video path for display
+
+### Gradio Gallery Format
+
+Input/output galleries use `type="filepath"` (string paths) rather than PIL Image tuples. Helper functions handle format compatibility for legacy tuple support.
+
+## Environment Variables
+
+No environment variables are required for basic operation. The application runs entirely with local models.
+
+## File Outputs
+
+Generated images are saved to `outputs/` directory with format:
+```
+output_{seed}_{index}_{timestamp_ms}.png
+```
+
+## Local Development and API Testing
+
+The `custom/` directory is fully gitignored and used for local development files. Specifically, it contains:
+
+- **API client scripts** - For testing the Gradio Space remotely via API after deployment to Hugging Face
+- **`API_GUIDE.txt`** - Auto-generated Gradio API documentation showing endpoint signatures and example usage
+- **Local testing environments** - Virtual environments or test data that shouldn't be committed
+
+**API Integration Pattern:**
+Once the Space is deployed to Hugging Face, you can interact with it programmatically using `gradio_client`:
+
+```python
+from gradio_client import Client, handle_file
+
+client = Client("clopeux/Pro-Realism-Edit-Studio")
+result = client.predict(
+    images=[],
+    prompt="Camera dollies forward, revealing more of the scene",
+    seed=42,
+    randomize_seed=False,
+    true_guidance_scale=1.0,
+    num_inference_steps=4,
+    height=1024,
+    width=1024,
+    api_name="/infer"
+)
+```
+
+The `custom/API_GUIDE.txt` contains full documentation of all available endpoints including `/infer`, `/turn_into_video`, and utility functions.
+
+## Development Notes
+
+- The model loads on startup and applies AOT compilation during first inference
+- Compilation uses dynamic shapes to support variable text/image sequence lengths
+- The transformer uses custom cache contexts ("cond"/"uncond") to optimize CFG passes
+- True CFG applies norm-based rescaling: `comb_pred * (cond_norm / noise_norm)`
+- FlashAttention 3 processor must be set before compilation

README.md

@@ -1,17 +1,17 @@
----
-title: Pro Realism Edit Studio
-emoji: 🎨
-colorFrom: indigo
-colorTo: blue
-sdk: gradio
-sdk_version: 5.49.1
-app_file: app.py
-pinned: false
-license: apache-2.0
-version: 0.1.3
-short_description: Powerful image editing - supports one or two input images.
----
-
-Pro Realism Edit Studio is a powerful image editor powered by [Qwen-Image-Edit-2511](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) with [Phr00t's Rapid-AIO v23](https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO) accelerated transformer for 4-step inference. Upload one or two input images, write a prompt, get high-quality results. Optional Enhance modes add lazy Nomos ATD upscaling, masked skin cleanup, and final detail grain.
-
-Video generation is disabled unless an owned Gradio video Space is configured with `VIDEO_SPACE_ID`.
+---
+title: Pro Realism Edit Studio
+emoji: 🎨
+colorFrom: indigo
+colorTo: blue
+sdk: gradio
+sdk_version: 5.49.1
+app_file: app.py
+pinned: false
+license: apache-2.0
+version: 0.1.3
+short_description: Powerful image editing - supports one or two input images.
+---
+
+Pro Realism Edit Studio is a powerful image editor powered by [Qwen-Image-Edit-2511](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) with [Phr00t's Rapid-AIO v23](https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO) accelerated transformer for 4-step inference. Upload one or two input images, write a prompt, get high-quality results. Optional Enhance modes add lazy Nomos ATD upscaling, masked skin cleanup, and final detail grain.
+
+Video generation is disabled unless an owned Gradio video Space is configured with `VIDEO_SPACE_ID`.

app.py

@@ -1,602 +1,986 @@
-import gradio as gr
-import numpy as np
-import random
-import torch
-import spaces
-
-from accelerate import init_empty_weights
-from collections import OrderedDict
-from PIL import Image, ImageEnhance, ImageFilter
-from diffusers.models import QwenImageTransformer2DModel as DiffusersQwenImageTransformer2DModel
-from diffusers.models.model_loading_utils import load_model_dict_into_meta
-from huggingface_hub import hf_hub_download
-from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
-from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
-from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
-from safetensors import safe_open
-
-import os
-import time  # Added for history update delay
-
-from gradio_client import Client, handle_file
-import tempfile
-
-BASE_MODEL_ID = "Qwen/Qwen-Image-Edit-2511"
-APP_VERSION = "0.1.3"
-PHR00T_REPO_ID = os.environ.get("PHR00T_REPO_ID", "Phr00t/Qwen-Image-Edit-Rapid-AIO").strip()
-RAPID_TRANSFORMER_FILENAME = os.environ.get(
-    "RAPID_TRANSFORMER_FILENAME",
-    "v23/Qwen-Rapid-AIO-NSFW-v23.safetensors",
-).strip()
-PHR00T_TRANSFORMER_PREFIX = "model.diffusion_model."
-VIDEO_SPACE_ID = os.environ.get("VIDEO_SPACE_ID", "").strip()
-UPSCALER_MODEL_ID = os.environ.get("UPSCALER_MODEL_ID", "Phips/4xNomos8k_atd_jpg").strip()
-UPSCALER_MODEL_FILENAME = os.environ.get("UPSCALER_MODEL_FILENAME", "4xNomos8k_atd_jpg.safetensors").strip()
-UPSCALER_TILE_SIZE = int(os.environ.get("UPSCALER_TILE_SIZE", "512"))
-UPSCALER_TILE_OVERLAP = int(os.environ.get("UPSCALER_TILE_OVERLAP", "32"))
-ENHANCE_MAX_INPUT_EDGE = int(os.environ.get("ENHANCE_MAX_INPUT_EDGE", "1280"))
-ENHANCE_GRAIN_STRENGTH = float(os.environ.get("ENHANCE_GRAIN_STRENGTH", "0.018"))
-
-ENHANCE_MODE_OFF = "Off"
-ENHANCE_MODE_UPSCALE = "Upscale"
-ENHANCE_MODE_CLEAN = "Clean"
-ENHANCE_MODE_MAX_DETAIL = "Max Detail"
-ENHANCE_MODE_CHOICES = [ENHANCE_MODE_OFF, ENHANCE_MODE_UPSCALE, ENHANCE_MODE_CLEAN, ENHANCE_MODE_MAX_DETAIL]
-
-_upscaler_model = None
-
-def turn_into_video(input_image, output_images, prompt, progress=gr.Progress(track_tqdm=True)):
-    if not VIDEO_SPACE_ID:
-        raise gr.Error("Video generation is not configured for this Space.")
-    if not input_image or not output_images:
-        raise gr.Error("Please generate an output image first.")
-
-    progress(0.02, desc="Preparing images...")
-
-    def extract_pil(img_entry):
-        if isinstance(img_entry, tuple) and isinstance(img_entry[0], Image.Image):
-            return img_entry[0]
-        elif isinstance(img_entry, Image.Image):
-            return img_entry
-        elif isinstance(img_entry, str):
-            return Image.open(img_entry)
-        else:
-            raise gr.Error(f"Unsupported image format: {type(img_entry)}")
-
-    start_img = extract_pil(input_image)
-    end_img   = extract_pil(output_images[0])
-
-    progress(0.10, desc="Saving temp files...")
-
-    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_start, \
-         tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_end:
-        start_img.save(tmp_start.name)
-        end_img.save(tmp_end.name)
-
-    progress(0.20, desc="Connecting to video Space...")
-
-    client = Client(VIDEO_SPACE_ID)
-
-    progress(0.35, desc="Generating video...")
-
-    video_path, seed = client.predict(
-        start_image_pil=handle_file(tmp_start.name),
-        end_image_pil=handle_file(tmp_end.name),
-        prompt=prompt or "smooth cinematic transition",
-        api_name="/generate_video"
-    )
-
-    progress(0.95, desc="Finalizing...")
-    print(video_path)
-    return video_path['video']
-
-
-def update_history(new_images, history):
-    """Updates the history gallery with the new images."""
-    time.sleep(0.5)  # Small delay to ensure images are ready
-    if history is None:
-        history = []
-    if new_images is not None and len(new_images) > 0:
-        if not isinstance(history, list):
-            history = list(history) if history else []
-        for img in new_images:
-            history.insert(0, img)
-    history = history[:20]  # Keep only last 20 images
-    return history
-
-def use_history_as_input(evt: gr.SelectData):
-    """Sets the selected history image into the Image 1 slot."""
-    if evt.value is not None:
-        # gr.Image with type='filepath' accepts a path directly.
-        return gr.update(value=evt.value)
-    return gr.update()
-
-# --- Model Loading ---
-dtype = torch.bfloat16
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-
-def load_phr00t_rapid_transformer(torch_dtype):
-    checkpoint_path = hf_hub_download(
-        repo_id=PHR00T_REPO_ID,
-        filename=RAPID_TRANSFORMER_FILENAME,
-    )
-    config = DiffusersQwenImageTransformer2DModel.load_config(
-        BASE_MODEL_ID,
-        subfolder="transformer",
-    )
-    with init_empty_weights():
-        transformer = DiffusersQwenImageTransformer2DModel.from_config(config)
-
-    expected_keys = set(transformer.state_dict().keys())
-    state_dict = OrderedDict()
-    with safe_open(checkpoint_path, framework="pt", device="cpu") as checkpoint:
-        for key in checkpoint.keys():
-            if not key.startswith(PHR00T_TRANSFORMER_PREFIX):
-                continue
-            mapped_key = key.removeprefix(PHR00T_TRANSFORMER_PREFIX)
-            if mapped_key in expected_keys:
-                state_dict[mapped_key] = checkpoint.get_tensor(key)
-
-    missing_keys = sorted(expected_keys.difference(state_dict.keys()))
-    if missing_keys:
-        sample = ", ".join(missing_keys[:20])
-        raise RuntimeError(
-            f"Phr00t Rapid-AIO transformer checkpoint is missing {len(missing_keys)} "
-            f"required diffusers keys after prefix conversion. First missing keys: {sample}"
-        )
-
-    load_model_dict_into_meta(transformer, state_dict, dtype=torch_dtype)
-    meta_parameters = [name for name, parameter in transformer.named_parameters() if parameter.is_meta]
-    if meta_parameters:
-        sample = ", ".join(meta_parameters[:20])
-        raise RuntimeError(
-            f"Phr00t Rapid-AIO transformer still has {len(meta_parameters)} meta parameters "
-            f"after loading. First meta parameters: {sample}"
-        )
-
-    transformer.eval()
-    return transformer
-
-
-# Load Qwen-Image-Edit-2511 with Phr00t's v23 accelerated transformer (4-step inference)
-pipe = QwenImageEditPlusPipeline.from_pretrained(
-    BASE_MODEL_ID,
-    transformer=load_phr00t_rapid_transformer(dtype),
-    torch_dtype=dtype
-).to(device)
-
-# Load next-scene LoRA for cinematic progression
-# Note: This LoRA was trained on 2509, may need testing with 2511/v23
-# TODO: Re-enable after testing base 2511/v23 works correctly
-# pipe.load_lora_weights(
-#     "lovis93/next-scene-qwen-image-lora-2509",
-#     weight_name="next-scene_lora-v2-3000.safetensors",
-#     adapter_name="next-scene"
-# )
-# pipe.set_adapters(["next-scene"], adapter_weights=[1.])
-# pipe.fuse_lora(adapter_names=["next-scene"], lora_scale=1.)
-# pipe.unload_lora_weights()
-
-
-# Apply the same optimizations from the first version
-pipe.transformer.__class__ = QwenImageTransformer2DModel
-pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
-
-# --- Ahead-of-time compilation ---
-# Note: optimize_pipeline_ handles text encoder offloading internally to save memory during torch.export
-# DISABLED 2026-05-12: HF build pipeline force-pins spaces==0.49.3 which has a regression in
-# zero.torch.patching._move() — NVML assert during worker_init kills AOTI compile at startup.
-# Restore once HF bumps the pipeline to spaces==0.50.0+.
-# optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")
-
-# --- UI Constants and Helpers ---
-MAX_SEED = np.iinfo(np.int32).max
-
-
-def load_upscaler_model():
-    global _upscaler_model
-    if _upscaler_model is not None:
-        return _upscaler_model
-
-    try:
-        import spandrel
-        import spandrel_extra_arches
-    except ImportError as exc:
-        raise gr.Error("Enhance mode requires spandrel and spandrel_extra_arches to be installed.") from exc
-
-    spandrel_extra_arches.install()
-    model_path = hf_hub_download(repo_id=UPSCALER_MODEL_ID, filename=UPSCALER_MODEL_FILENAME)
-    model = spandrel.ModelLoader().load_from_file(model_path)
-    model.eval().to(device)
-    _upscaler_model = model
-    return _upscaler_model
-
-
-def image_to_tensor(image):
-    array = np.asarray(image.convert("RGB")).astype(np.float32) / 255.0
-    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
-    return tensor.to(device)
-
-
-def tensor_to_image(tensor):
-    array = tensor.squeeze(0).detach().float().cpu().clamp(0, 1).permute(1, 2, 0).numpy()
-    return Image.fromarray((array * 255.0).round().astype(np.uint8), mode="RGB")
-
-
-def validate_enhance_input_size(image):
-    max_edge = max(image.size)
-    if max_edge > ENHANCE_MAX_INPUT_EDGE:
-        raise gr.Error(
-            f"Enhance mode accepts images up to {ENHANCE_MAX_INPUT_EDGE}px on the longest edge. "
-            f"Current image is {image.width}x{image.height}."
-        )
-
-
-def tile_upscale(image):
-    validate_enhance_input_size(image)
-    model = load_upscaler_model()
-    tensor = image_to_tensor(image)
-    _, _, height, width = tensor.shape
-    tile_size = max(64, min(UPSCALER_TILE_SIZE, height, width))
-    overlap = max(0, min(UPSCALER_TILE_OVERLAP, tile_size // 2))
-    step = max(1, tile_size - overlap)
-    y_positions = sorted(set(list(range(0, height, step)) + [max(0, height - tile_size)]))
-    x_positions = sorted(set(list(range(0, width, step)) + [max(0, width - tile_size)]))
-    output = None
-    weights = None
-
-    with torch.inference_mode():
-        for y in y_positions:
-            for x in x_positions:
-                y1 = min(y + tile_size, height)
-                x1 = min(x + tile_size, width)
-                tile = tensor[:, :, y:y1, x:x1]
-                upscaled_tile = model(tile).clamp(0, 1)
-                scale_y = upscaled_tile.shape[-2] // tile.shape[-2]
-                scale_x = upscaled_tile.shape[-1] // tile.shape[-1]
-                if output is None:
-                    output = torch.zeros(
-                        (1, 3, height * scale_y, width * scale_x),
-                        dtype=upscaled_tile.dtype,
-                        device=upscaled_tile.device,
-                    )
-                    weights = torch.zeros_like(output)
-                oy0, oy1 = y * scale_y, y1 * scale_y
-                ox0, ox1 = x * scale_x, x1 * scale_x
-                output[:, :, oy0:oy1, ox0:ox1] += upscaled_tile
-                weights[:, :, oy0:oy1, ox0:ox1] += 1
-
-    output = output / weights.clamp_min(1)
-    return tensor_to_image(output)
-
-
-def skin_repair_mask(image):
-    ycbcr = np.asarray(image.convert("YCbCr"))
-    cb = ycbcr[:, :, 1]
-    cr = ycbcr[:, :, 2]
-    mask = (
-        (cr >= 135)
-        & (cr <= 180)
-        & (cb >= 75)
-        & (cb <= 135)
-    ).astype(np.uint8) * 255
-    mask_image = Image.fromarray(mask, mode="L")
-    return mask_image.filter(ImageFilter.GaussianBlur(radius=1.2))
-
-
-def repair_skin_texture(image):
-    base = image.convert("RGB")
-    mask = skin_repair_mask(base)
-    repaired = base.filter(ImageFilter.MedianFilter(size=3)).filter(ImageFilter.GaussianBlur(radius=0.35))
-    blended = Image.composite(repaired, base, mask)
-    return ImageEnhance.Sharpness(blended).enhance(1.08)
-
-
-def add_film_grain(image, seed):
-    base = image.convert("RGB")
-    array = np.asarray(base).astype(np.float32)
-    rng = np.random.default_rng(seed)
-    grain = rng.normal(0.0, 255.0 * ENHANCE_GRAIN_STRENGTH, size=(array.shape[0], array.shape[1], 1))
-    array = np.clip(array + grain, 0, 255)
-    return Image.fromarray(array.astype(np.uint8), mode="RGB")
-
-
-def apply_enhancement(image, enhance_mode, seed=0, progress=None):
-    mode = enhance_mode or ENHANCE_MODE_OFF
-    if mode not in ENHANCE_MODE_CHOICES:
-        raise gr.Error(f"Unknown enhance mode: {mode}")
-    if mode == ENHANCE_MODE_OFF:
-        return image
-
-    enhanced = image.convert("RGB")
-    if mode in (ENHANCE_MODE_CLEAN, ENHANCE_MODE_MAX_DETAIL):
-        if progress:
-            progress(0.76, desc="Repairing skin texture...")
-        enhanced = repair_skin_texture(enhanced)
-
-    if mode in (ENHANCE_MODE_UPSCALE, ENHANCE_MODE_CLEAN, ENHANCE_MODE_MAX_DETAIL):
-        if progress:
-            progress(0.82, desc="Upscaling image...")
-        enhanced = tile_upscale(enhanced)
-
-    if mode == ENHANCE_MODE_MAX_DETAIL:
-        if progress:
-            progress(0.92, desc="Adding final grain and sharpness...")
-        enhanced = ImageEnhance.Sharpness(enhanced).enhance(1.12)
-        enhanced = add_film_grain(enhanced, seed)
-
-    return enhanced
-
-def use_output_as_input(output_images):
-    """Move the first output image into the Image 1 slot."""
-    if not output_images:
-        return gr.update()
-    first = output_images[0]
-    # Gallery items can be filepath strings or (filepath, label) tuples.
-    path = first[0] if isinstance(first, (list, tuple)) else first
-    return gr.update(value=path)
-
-# --- Main Inference Function (with hardcoded negative prompt) ---
-@spaces.GPU(duration=60)
-def infer(
-    image_1,
-    image_2,
-    prompt,
-    seed=42,
-    randomize_seed=False,
-    true_guidance_scale=1.0,
-    num_inference_steps=4,
-    height=None,
-    width=None,
-    enhance_mode=ENHANCE_MODE_OFF,
-    num_images_per_prompt=1,
-    progress=gr.Progress(track_tqdm=True),
-):
-    """
-    Generates an image using the local Qwen-Image diffusers pipeline.
-    """
-    # Hardcode the negative prompt as requested
-    negative_prompt = " "
-
-    if randomize_seed:
-        seed = random.randint(0, MAX_SEED)
-
-    # Set up the generator for reproducibility
-    generator = torch.Generator(device=device).manual_seed(seed)
-
-    # Load input images into PIL Images — two optional slots.
-    pil_images = []
-    for img in (image_1, image_2):
-        if img is None:
-            continue
-        try:
-            if isinstance(img, str):
-                pil_images.append(Image.open(img).convert("RGB"))
-            elif isinstance(img, Image.Image):
-                pil_images.append(img.convert("RGB"))
-            elif hasattr(img, "name"):
-                pil_images.append(Image.open(img.name).convert("RGB"))
-        except Exception:
-            continue
-
-    if height==256 and width==256:
-        height, width = None, None
-    print(f"Calling pipeline with prompt: '{prompt}'")
-    print(f"Negative Prompt: '{negative_prompt}'")
-    print(f"Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}, Size: {width}x{height}")
-
-    # Generate the image
-    images_pil = pipe(
-        image=pil_images if len(pil_images) > 0 else None,
-        prompt=prompt,
-        height=height,
-        width=width,
-        negative_prompt=negative_prompt,
-        num_inference_steps=num_inference_steps,
-        generator=generator,
-        true_cfg_scale=true_guidance_scale,
-        num_images_per_prompt=num_images_per_prompt,
-    ).images
-
-    if enhance_mode != ENHANCE_MODE_OFF:
-        images_pil = [
-            apply_enhancement(img, enhance_mode, seed=seed + idx, progress=progress)
-            for idx, img in enumerate(images_pil)
-        ]
-
-    # Save images to temporary files for proper serving
-    output_paths = []
-    os.makedirs("outputs", exist_ok=True)
-    for idx, img in enumerate(images_pil):
-        output_path = f"outputs/output_{seed}_{idx}_{int(time.time()*1000)}.png"
-        img.save(output_path)
-        output_paths.append(output_path)
-
-    # Return image paths, seed, and make buttons visible when their feature is configured.
-    return output_paths, seed, gr.update(visible=True), gr.update(visible=bool(VIDEO_SPACE_ID))
-
-
-# --- UI Layout ---
-css = """
-#col-container {
-    margin: 0 auto;
-    max-width: 1024px;
-}
-#logo-title {
-    text-align: center;
-}
-#logo-title h1 {
-    margin-bottom: 0;
-}
-#logo-title h2 {
-    color: #5b47d1;
-    font-style: italic;
-    margin-top: 0;
-}
-#edit_text{margin-top: -62px !important}
-"""
-
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.HTML(f"""
-        <!-- v{APP_VERSION} -->
-        <div id="logo-title">
-            <h1>Pro Realism Edit Studio</h1>
-            <h2>Rapid Edit ⚡</h2>
-        </div>
-        """)
-        gr.Markdown("""
-        This demo uses [Qwen-Image-Edit-2511](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) with [Phr00t's Rapid-AIO v23](https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO) accelerated transformer + [AoT compilation & FA3](https://huggingface.co/blog/zerogpu-aoti) for fast 4-step inference.
-
-        Upload an image and enter your prompt to edit it. The model will use your prompt exactly as provided.
-        """)
-        with gr.Row():
-            with gr.Column():
-                with gr.Row():
-                    image_1 = gr.Image(label="Image 1", type="filepath", interactive=True)
-                    image_2 = gr.Image(label="Image 2 (optional)", type="filepath", interactive=True)
-
-                prompt = gr.Text(
-                    label="Prompt 🪄",
-                    show_label=True,
-                    placeholder="Enter your prompt here...",
-            )
-                enhance_mode = gr.Radio(
-                    label="Enhance / Amelioration",
-                    choices=ENHANCE_MODE_CHOICES,
-                    value=ENHANCE_MODE_OFF,
-                    interactive=True,
-                )
-                run_button = gr.Button("Edit!", variant="primary")
-                
-                with gr.Accordion("Advanced Settings", open=False):
-                    
-        
-                    seed = gr.Slider(
-                        label="Seed",
-                        minimum=0,
-                        maximum=MAX_SEED,
-                        step=1,
-                        value=0,
-                    )
-        
-                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-        
-                    with gr.Row():
-        
-                        true_guidance_scale = gr.Slider(
-                            label="True guidance scale",
-                            minimum=1.0,
-                            maximum=10.0,
-                            step=0.1,
-                            value=1.0
-                        )
-
-                        num_inference_steps = gr.Slider(
-                            label="Number of inference steps",
-                            minimum=1,
-                            maximum=40,
-                            step=1,
-                            value=4,
-                        )
-                        
-                        height = gr.Slider(
-                            label="Height",
-                            minimum=256,
-                            maximum=2048,
-                            step=8,
-                            value=None,
-                        )
-                        
-                        width = gr.Slider(
-                            label="Width",
-                            minimum=256,
-                            maximum=2048,
-                            step=8,
-                            value=None,
-                        )
-
-
-
-            with gr.Column():
-                result = gr.Gallery(label="Result", show_label=False, type="filepath")
-                with gr.Row():
-                    use_output_btn = gr.Button("↗️ Use as input", variant="secondary", size="sm", visible=False)
-                    turn_video_btn = gr.Button("🎬 Turn into Video", variant="secondary", size="sm", visible=False)
-                output_video = gr.Video(label="Generated Video", autoplay=True, visible=False)
-
-                with gr.Row(visible=False):
-                    gr.Markdown("### 📜 History")
-                    clear_history_button = gr.Button("🗑️ Clear History", size="sm", variant="stop")
-                
-                history_gallery = gr.Gallery(
-                    label="Click any image to use as input", 
-                    interactive=False,
-                    show_label=True,
-                    visible=False
-                )
-
-
-
-        
-
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            image_1,
-            image_2,
-            prompt,
-            seed,
-            randomize_seed,
-            true_guidance_scale,
-            num_inference_steps,
-            height,
-            width,
-            enhance_mode,
-        ],
-        outputs=[result, seed, use_output_btn, turn_video_btn],
-
-    ).then(
-    fn=update_history,
-    inputs=[result, history_gallery],
-    outputs=history_gallery,
-
-    )
-
-    # Add the new event handler for the "Use Output as Input" button
-    use_output_btn.click(
-        fn=use_output_as_input,
-        inputs=[result],
-        outputs=[image_1]
-    )
-
-    # History gallery event handlers
-    history_gallery.select(
-        fn=use_history_as_input,
-        inputs=None,
-        outputs=[image_1],
-
-    )
-    
-    clear_history_button.click(
-        fn=lambda: [],
-        inputs=None,
-        outputs=history_gallery,
-
-    )
-
-    turn_video_btn.click(
-    fn=lambda: gr.update(visible=True),   
-    inputs=None,
-    outputs=[output_video],
-).then(
-    fn=turn_into_video,
-    inputs=[image_1, result, prompt],
-    outputs=[output_video],
-)
-
-
-if __name__ == "__main__":
-    demo.launch()
+#!/usr/bin/env python3
+"""
+Pro Realism Edit Studio - Enhanced Edition
+=========================================
+
+Advanced image editing and enhancement studio powered by:
+- Qwen-Image-Edit-2511 with Phr00t's Rapid-AIO v23 accelerated transformer
+- Real-ESRGAN for high-quality upscaling
+- GFPGAN/CodeFormer for face restoration
+- Multi-stage detail enhancement pipeline
+
+Author: Enhanced with Hugging Face CLI and image generation expertise
+Version: 1.0.0
+"""
+
+import gradio as gr
+import numpy as np
+import random
+import torch
+import spaces
+import os
+import time
+import tempfile
+from pathlib import Path
+
+# Advanced imports
+from accelerate import init_empty_weights
+from collections import OrderedDict
+from PIL import Image, ImageEnhance, ImageFilter, ImageOps
+from diffusers.models import QwenImageTransformer2DModel as DiffusersQwenImageTransformer2DModel
+from diffusers.models.model_loading_utils import load_model_dict_into_meta
+from huggingface_hub import hf_hub_download, HfApi, login, whoami
+from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
+from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
+from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
+from safetensors import safe_open
+
+from gradio_client import Client, handle_file
+
+# ============================================================================
+# CONFIGURATION - Model IDs and Parameters
+# ============================================================================
+
+# Base model configuration
+BASE_MODEL_ID = "Qwen/Qwen-Image-Edit-2511"
+APP_VERSION = "1.0.0"
+PHR00T_REPO_ID = os.environ.get("PHR00T_REPO_ID", "Phr00t/Qwen-Image-Edit-Rapid-AIO").strip()
+RAPID_TRANSFORMER_FILENAME = os.environ.get(
+    "RAPID_TRANSFORMER_FILENAME",
+    "v23/Qwen-Rapid-AIO-NSFW-v23.safetensors",
+).strip()
+PHR00T_TRANSFORMER_PREFIX = "model.diffusion_model."
+VIDEO_SPACE_ID = os.environ.get("VIDEO_SPACE_ID", "").strip()
+
+# Enhanced Upscaler Configuration
+UPSCALER_MODEL_ID = os.environ.get("UPSCALER_MODEL_ID", "ai-forever/Real-ESRGAN").strip()
+UPSCALER_MODEL_FILENAME = os.environ.get("UPSCALER_MODEL_FILENAME", "RealESRGAN_x4plus.pth").strip()
+UPSCALER_TILE_SIZE = int(os.environ.get("UPSCALER_TILE_SIZE", "512"))
+UPSCALER_TILE_OVERLAP = int(os.environ.get("UPSCALER_TILE_OVERLAP", "64"))
+ENHANCE_MAX_INPUT_EDGE = int(os.environ.get("ENHANCE_MAX_INPUT_EDGE", "2048"))
+ENHANCE_GRAIN_STRENGTH = float(os.environ.get("ENHANCE_GRAIN_STRENGTH", "0.015"))
+
+# Face Restoration Configuration
+FACE_RESTORATION_MODEL = os.environ.get("FACE_RESTORATION_MODEL", "Xintao/GFPGAN").strip()
+FACE_RESTORATION_WEIGHTS = os.environ.get("FACE_RESTORATION_WEIGHTS", "GFPGANv1.3.pth").strip()
+
+# Advanced Detail Enhancement Configuration
+DETAIL_ENHANCEMENT_ENABLED = os.environ.get("DETAIL_ENHANCEMENT_ENABLED", "true").lower() == "true"
+SMART_SHARPENING_STRENGTH = float(os.environ.get("SMART_SHARPENING_STRENGTH", "1.15"))
+
+# ============================================================================
+# ENHANCEMENT MODES
+# ============================================================================
+
+ENHANCE_MODE_OFF = "Off"
+ENHANCE_MODE_UPSCALE = "Upscale Only"
+ENHANCE_MODE_CLEAN = "Clean & Restore"
+ENHANCE_MODE_MAX_DETAIL = "Max Detail"
+ENHANCE_MODE_FACE_ENHANCE = "Face Enhance"
+ENHANCE_MODE_FULL_ENHANCE = "Full Enhance"
+ENHANCE_MODE_CHOICES = [
+    ENHANCE_MODE_OFF, 
+    ENHANCE_MODE_UPSCALE, 
+    ENHANCE_MODE_CLEAN, 
+    ENHANCE_MODE_MAX_DETAIL,
+    ENHANCE_MODE_FACE_ENHANCE,
+    ENHANCE_MODE_FULL_ENHANCE
+]
+
+# ============================================================================
+# GLOBAL MODEL CACHE
+# ============================================================================
+
+_upscaler_model = None
+_face_restoration_model = None
+_detail_enhancement_model = None
+
+# ============================================================================
+# HUGGING FACE CLI EXPERT FUNCTIONS
+# ============================================================================
+
+def download_model_with_retry(repo_id, filename, max_retries=3):
+    """Download model with retry logic and error handling"""
+    for attempt in range(max_retries):
+        try:
+            return hf_hub_download(repo_id=repo_id, filename=filename)
+        except Exception as e:
+            if attempt == max_retries - 1:
+                raise RuntimeError(f"Failed to download {filename} from {repo_id} after {max_retries} attempts: {e}")
+            time.sleep(2 ** attempt)
+    return None
+
+# ============================================================================
+# VIDEO GENERATION (Preserved from original)
+# ============================================================================
+
+def turn_into_video(input_image, output_images, prompt, progress=gr.Progress(track_tqdm=True)):
+    """Convert image edit into video transition"""
+    if not VIDEO_SPACE_ID:
+        raise gr.Error("Video generation is not configured for this Space.")
+    if not input_image or not output_images:
+        raise gr.Error("Please generate an output image first.")
+
+    progress(0.02, desc="Preparing images...")
+
+    def extract_pil(img_entry):
+        if isinstance(img_entry, tuple) and isinstance(img_entry[0], Image.Image):
+            return img_entry[0]
+        elif isinstance(img_entry, Image.Image):
+            return img_entry
+        elif isinstance(img_entry, str):
+            return Image.open(img_entry)
+        else:
+            raise gr.Error(f"Unsupported image format: {type(img_entry)}")
+
+    start_img = extract_pil(input_image)
+    end_img = extract_pil(output_images[0])
+
+    progress(0.10, desc="Saving temp files...")
+
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_start, \
+         tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_end:
+        start_img.save(tmp_start.name)
+        end_img.save(tmp_end.name)
+
+    progress(0.20, desc="Connecting to video Space...")
+
+    client = Client(VIDEO_SPACE_ID)
+
+    progress(0.35, desc="Generating video...")
+
+    video_path, seed = client.predict(
+        start_image_pil=handle_file(tmp_start.name),
+        end_image_pil=handle_file(tmp_end.name),
+        prompt=prompt or "smooth cinematic transition",
+        api_name="/generate_video"
+    )
+
+    progress(0.95, desc="Finalizing...")
+    return video_path['video']
+
+
+# ============================================================================
+# HISTORY MANAGEMENT (Enhanced)
+# ============================================================================
+
+def update_history(new_images, history):
+    """Updates the history gallery with the new images."""
+    time.sleep(0.3)
+    if history is None:
+        history = []
+    if new_images is not None and len(new_images) > 0:
+        if not isinstance(history, list):
+            history = list(history) if history else []
+        for img in new_images:
+            history.insert(0, img)
+    history = history[:50]
+    return history
+
+def use_history_as_input(evt: gr.SelectData):
+    """Sets the selected history image into the Image 1 slot."""
+    if evt.value is not None:
+        return gr.update(value=evt.value)
+    return gr.update()
+
+# ============================================================================
+# MODEL LOADING (Enhanced with better error handling)
+# ============================================================================
+
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def load_phr00t_rapid_transformer(torch_dtype):
+    """Load Phr00t's Rapid-AIO v23 transformer with enhanced error handling"""
+    checkpoint_path = download_model_with_retry(PHR00T_REPO_ID, RAPID_TRANSFORMER_FILENAME)
+    
+    try:
+        config = DiffusersQwenImageTransformer2DModel.load_config(
+            BASE_MODEL_ID,
+            subfolder="transformer",
+        )
+    except Exception as e:
+        raise RuntimeError(f"Failed to load config for {BASE_MODEL_ID}: {e}")
+    
+    with init_empty_weights():
+        transformer = DiffusersQwenImageTransformer2DModel.from_config(config)
+
+    expected_keys = set(transformer.state_dict().keys())
+    state_dict = OrderedDict()
+    
+    try:
+        with safe_open(checkpoint_path, framework="pt", device="cpu") as checkpoint:
+            for key in checkpoint.keys():
+                if not key.startswith(PHR00T_TRANSFORMER_PREFIX):
+                    continue
+                mapped_key = key.removeprefix(PHR00T_TRANSFORMER_PREFIX)
+                if mapped_key in expected_keys:
+                    state_dict[mapped_key] = checkpoint.get_tensor(key)
+    except Exception as e:
+        raise RuntimeError(f"Failed to load checkpoint from {checkpoint_path}: {e}")
+
+    missing_keys = sorted(expected_keys.difference(state_dict.keys()))
+    if missing_keys:
+        sample = ", ".join(missing_keys[:20])
+        raise RuntimeError(
+            f"Phr00t Rapid-AIO transformer checkpoint is missing {len(missing_keys)} "
+            f"required diffusers keys after prefix conversion. First missing keys: {sample}"
+        )
+
+    try:
+        load_model_dict_into_meta(transformer, state_dict, dtype=torch_dtype)
+    except Exception as e:
+        raise RuntimeError(f"Failed to load state dict into meta: {e}")
+    
+    meta_parameters = [name for name, parameter in transformer.named_parameters() if parameter.is_meta]
+    if meta_parameters:
+        sample = ", ".join(meta_parameters[:20])
+        raise RuntimeError(
+            f"Phr00t Rapid-AIO transformer still has {len(meta_parameters)} meta parameters "
+            f"after loading. First meta parameters: {sample}"
+        )
+
+    transformer.eval()
+    return transformer
+
+# Load main pipeline
+try:
+    pipe = QwenImageEditPlusPipeline.from_pretrained(
+        BASE_MODEL_ID,
+        transformer=load_phr00t_rapid_transformer(dtype),
+        torch_dtype=dtype
+    ).to(device)
+    print("Successfully loaded Qwen-Image-Edit-2511 with Rapid-AIO v23 transformer")
+except Exception as e:
+    print(f"Failed to load main pipeline: {e}")
+    raise
+
+# Apply optimizations
+pipe.transformer.__class__ = QwenImageTransformer2DModel
+pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+print("Applied FA3 attention processor optimization")
+
+# ============================================================================
+# ENHANCED UPSCALER (Real-ESRGAN based)
+# ============================================================================
+
+def load_upscaler_model():
+    """Load Real-ESRGAN model for high-quality upscaling"""
+    global _upscaler_model
+    if _upscaler_model is not None:
+        return _upscaler_model
+
+    try:
+        import spandrel
+        import spandrel_extra_arches
+        spandrel_extra_arches.install()
+    except ImportError as exc:
+        raise gr.Error("Enhance mode requires spandrel and spandrel_extra_arches to be installed. "
+                       "Install with: pip install spandrel spandrel_extra_arches") from exc
+
+    try:
+        model_path = download_model_with_retry(UPSCALER_MODEL_ID, UPSCALER_MODEL_FILENAME)
+        model = spandrel.ModelLoader().load_from_file(model_path)
+        model.eval().to(device)
+        _upscaler_model = model
+        print(f"Successfully loaded upscaler: {UPSCALER_MODEL_ID}/{UPSCALER_MODEL_FILENAME}")
+        return _upscaler_model
+    except Exception as e:
+        print(f"Failed to load upscaler model: {e}")
+        print("Falling back to Nomos upscaler...")
+        try:
+            model_path = download_model_with_retry("Phips/4xNomos8k_atd_jpg", "4xNomos8k_atd_jpg.safetensors")
+            model = spandrel.ModelLoader().load_from_file(model_path)
+            model.eval().to(device)
+            _upscaler_model = model
+            return _upscaler_model
+        except Exception as fallback_error:
+            raise gr.Error(f"Failed to load all upscaler models: {e} | {fallback_error}")
+
+def image_to_tensor(image):
+    """Convert PIL Image to tensor"""
+    array = np.asarray(image.convert("RGB")).astype(np.float32) / 255.0
+    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
+    return tensor.to(device)
+
+def tensor_to_image(tensor):
+    """Convert tensor to PIL Image"""
+    array = tensor.squeeze(0).detach().float().cpu().clamp(0, 1).permute(1, 2, 0).numpy()
+    return Image.fromarray((array * 255.0).round().astype(np.uint8), mode="RGB")
+
+def validate_enhance_input_size(image):
+    """Validate image size for enhancement"""
+    max_edge = max(image.size)
+    if max_edge > ENHANCE_MAX_INPUT_EDGE:
+        raise gr.Error(
+            f"Enhance mode accepts images up to {ENHANCE_MAX_INPUT_EDGE}px on the longest edge. "
+            f"Current image is {image.width}x{image.height}. "
+            f"Consider resizing your image first."
+        )
+
+def advanced_tile_upscale(image, scale=4):
+    """
+    Advanced tiling upscaler with improved blending and edge handling
+    Uses Real-ESRGAN for superior quality compared to Nomos
+    """
+    validate_enhance_input_size(image)
+    model = load_upscaler_model()
+    tensor = image_to_tensor(image)
+    _, _, height, width = tensor.shape
+    
+    base_tile_size = UPSCALER_TILE_SIZE
+    optimal_tile_size = min(base_tile_size, max(height, width) // 2)
+    tile_size = max(64, optimal_tile_size)
+    overlap = max(0, min(UPSCALER_TILE_OVERLAP, tile_size // 2))
+    step = max(1, tile_size - overlap)
+    
+    y_positions = list(range(0, height, step))
+    if y_positions[-1] + tile_size < height:
+        y_positions.append(max(0, height - tile_size))
+    
+    x_positions = list(range(0, width, step))
+    if x_positions[-1] + tile_size < width:
+        x_positions.append(max(0, width - tile_size))
+    
+    y_positions = sorted(set(y_positions))
+    x_positions = sorted(set(x_positions))
+    
+    output = None
+    weights = None
+
+    with torch.inference_mode():
+        for y in y_positions:
+            for x in x_positions:
+                y1 = min(y + tile_size, height)
+                x1 = min(x + tile_size, width)
+                tile = tensor[:, :, y:y1, x:x1]
+                
+                upscaled_tile = model(tile).clamp(0, 1)
+                
+                scale_y = upscaled_tile.shape[-2] // tile.shape[-2]
+                scale_x = upscaled_tile.shape[-1] // tile.shape[-1]
+                
+                if output is None:
+                    output = torch.zeros(
+                        (1, 3, height * scale_y, width * scale_x),
+                        dtype=upscaled_tile.dtype,
+                        device=upscaled_tile.device,
+                    )
+                    weights = torch.zeros_like(output)
+                
+                oy0, oy1 = y * scale_y, y1 * scale_y
+                ox0, ox1 = x * scale_x, x1 * scale_x
+                output[:, :, oy0:oy1, ox0:ox1] += upscaled_tile
+                weights[:, :, oy0:oy1, ox0:ox1] += 1
+
+    output = output / weights.clamp_min(1)
+    return tensor_to_image(output)
+
+# ============================================================================
+# ENHANCED DETAILER (Multi-stage processing)
+# ============================================================================
+
+def smart_sharpen(image, strength=1.15):
+    """
+    Smart sharpening with edge detection to avoid oversharpening smooth areas
+    """
+    if strength <= 0:
+        return image
+    
+    img_array = np.array(image.convert("RGB"))
+    
+    if strength > 1.0:
+        enhanced = ImageEnhance.Sharpness(image).enhance(strength)
+        
+        gray = image.convert("L")
+        edges = gray.filter(ImageFilter.FIND_EDGES)
+        edge_mask = edges.filter(ImageFilter.GaussianBlur(radius=1))
+        edge_mask = edge_mask.point(lambda x: min(x * 0.3, 255))
+        
+        sharpened_array = np.array(enhanced)
+        original_array = img_array
+        edge_array = np.array(edge_mask).astype(float) / 255.0
+        
+        for c in range(3):
+            sharpened_array[:, :, c] = (
+                edge_array * sharpened_array[:, :, c] + 
+                (1 - edge_array) * original_array[:, :, c]
+            )
+        
+        image = Image.fromarray(np.clip(sharpened_array, 0, 255).astype(np.uint8))
+    
+    return image
+
+def add_ultra_detail(image, strength=0.8):
+    """
+    Add ultra-fine details using high-frequency enhancement
+    """
+    if strength <= 0:
+        return image
+    
+    original = image.convert("RGB")
+    blurred = original.filter(ImageFilter.GaussianBlur(radius=2))
+    
+    high_freq = ImageChops.subtract(original, blurred)
+    high_freq_enhanced = ImageEnhance.Contrast(high_freq).enhance(1.0 + strength)
+    result = ImageChops.add(original, high_freq_enhanced)
+    
+    return result
+
+def apply_high_frequency_details(image, amount=0.6):
+    """
+    Apply high-frequency detail enhancement for crisp textures
+    """
+    if amount <= 0:
+        return image
+    
+    scales = [1, 2, 4]
+    result = image.convert("RGB")
+    
+    for scale in scales:
+        blurred = result.filter(ImageFilter.GaussianBlur(radius=scale))
+        high_freq = ImageChops.subtract(result, blurred)
+        enhanced_hf = ImageEnhance.Contrast(high_freq).enhance(1.0 + amount * 0.3)
+        result = ImageChops.add(result, enhanced_hf)
+    
+    return result
+
+# ============================================================================
+# ENHANCED CLEANER (Face Restoration + Artifact Removal)
+# ============================================================================
+
+def load_face_restoration_model():
+    """Load GFPGAN model for face restoration"""
+    global _face_restoration_model
+    if _face_restoration_model is not None:
+        return _face_restoration_model
+    
+    try:
+        from gfpgan import GFPGANer
+        model_path = download_model_with_retry(FACE_RESTORATION_MODEL, FACE_RESTORATION_WEIGHTS)
+        
+        restorer = GFPGANer(
+            model_path=model_path,
+            upscale=1,
+            arch='clean',
+            channel_multiplier=2,
+            bg_upsampler=None
+        )
+        
+        _face_restoration_model = restorer
+        print("Successfully loaded GFPGAN face restoration model")
+        return _face_restoration_model
+        
+    except ImportError:
+        print("GFPGAN not available, face restoration will use fallback methods")
+        return None
+    except Exception as e:
+        print(f"Failed to load face restoration model: {e}")
+        return None
+
+def restore_faces(image):
+    """Restore faces in an image using GFPGAN"""
+    restorer = load_face_restoration_model()
+    if restorer is None:
+        print("Face restoration model not available, using skin repair fallback")
+        return repair_skin_texture(image)
+    
+    try:
+        img_array = np.array(image.convert("RGB"))
+        restored_array, _ = restorer.enhance(img_array, has_aligned=False, only_center_face=False, paste_back=True)
+        restored_image = Image.fromarray(restored_array.astype(np.uint8))
+        return restored_image
+    except Exception as e:
+        print(f"Face restoration failed: {e}, using skin repair fallback")
+        return repair_skin_texture(image)
+
+def remove_artifacts(image):
+    """Remove compression artifacts and noise"""
+    denoised = image.filter(ImageFilter.MedianFilter(size=3))
+    smoothed = denoised.filter(ImageFilter.GaussianBlur(radius=0.5))
+    result = Image.blend(image, smoothed, alpha=0.3)
+    return result
+
+def enhanced_skin_repair(image):
+    """Enhanced skin repair with better color detection and blending"""
+    base = image.convert("RGB")
+    
+    ycbcr = np.asarray(base.convert("YCbCr"))
+    y, cb, cr = ycbcr[:, :, 0], ycbcr[:, :, 1], ycbcr[:, :, 2]
+    
+    skin_mask = (
+        (cr > 130) & (cr < 170) &
+        (cb > 70) & (cb < 140) &
+        (y > 80)
+    ).astype(np.uint8) * 255
+    
+    try:
+        import cv2
+        kernel = np.ones((5, 5), np.uint8)
+        skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_OPEN, kernel)
+        skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_CLOSE, kernel)
+        skin_mask = cv2.GaussianBlur(skin_mask, (7, 7), 0)
+    except ImportError:
+        from scipy import ndimage
+        skin_mask = ndimage.binary_opening(skin_mask > 128, structure=np.ones((3, 3))).astype(np.uint8) * 255
+        skin_mask = ndimage.gaussian_filter(skin_mask, sigma=3)
+    
+    mask_image = Image.fromarray(skin_mask, mode="L")
+    repaired = base.filter(ImageFilter.MedianFilter(size=3))
+    repaired = repaired.filter(ImageFilter.GaussianBlur(radius=0.4))
+    non_skin = ImageOps.invert(mask_image)
+    sharpened = ImageEnhance.Sharpness(base).enhance(1.15)
+    blended = Image.composite(repaired, sharpened, mask_image)
+    result = ImageEnhance.Sharpness(blended).enhance(1.05)
+    return result
+
+# Original skin repair functions (preserved for compatibility)
+def skin_repair_mask(image):
+    ycbcr = np.asarray(image.convert("YCbCr"))
+    cb = ycbcr[:, :, 1]
+    cr = ycbcr[:, :, 2]
+    mask = (
+        (cr >= 135)
+        & (cr <= 180)
+        & (cb >= 75)
+        & (cb <= 135)
+    ).astype(np.uint8) * 255
+    mask_image = Image.fromarray(mask, mode="L")
+    return mask_image.filter(ImageFilter.GaussianBlur(radius=1.2))
+
+def repair_skin_texture(image):
+    base = image.convert("RGB")
+    mask = skin_repair_mask(base)
+    repaired = base.filter(ImageFilter.MedianFilter(size=3)).filter(ImageFilter.GaussianBlur(radius=0.35))
+    blended = Image.composite(repaired, base, mask)
+    return ImageEnhance.Sharpness(blended).enhance(1.08)
+
+def add_film_grain(image, seed):
+    base = image.convert("RGB")
+    array = np.asarray(base).astype(np.float32)
+    rng = np.random.default_rng(seed)
+    grain = rng.normal(0.0, 255.0 * ENHANCE_GRAIN_STRENGTH, size=(array.shape[0], array.shape[1], 1))
+    array = np.clip(array + grain, 0, 255)
+    return Image.fromarray(array.astype(np.uint8), mode="RGB")
+
+# ============================================================================
+# ENHANCED APPLY ENHANCEMENT (Main enhancement pipeline)
+# ============================================================================
+
+def apply_enhancement(image, enhance_mode, seed=0, progress=None):
+    """
+    Apply various enhancement modes to the image
+    """
+    mode = enhance_mode or ENHANCE_MODE_OFF
+    if mode not in ENHANCE_MODE_CHOICES:
+        raise gr.Error(f"Unknown enhance mode: {mode}")
+    if mode == ENHANCE_MODE_OFF:
+        return image
+
+    enhanced = image.convert("RGB")
+    
+    total_steps = 0
+    if mode == ENHANCE_MODE_UPSCALE:
+        total_steps = 1
+    elif mode == ENHANCE_MODE_CLEAN:
+        total_steps = 3
+    elif mode == ENHANCE_MODE_MAX_DETAIL:
+        total_steps = 4
+    elif mode == ENHANCE_MODE_FACE_ENHANCE:
+        total_steps = 2
+    elif mode == ENHANCE_MODE_FULL_ENHANCE:
+        total_steps = 5
+    
+    step = 0
+    
+    if mode == ENHANCE_MODE_FACE_ENHANCE:
+        if progress:
+            step += 1
+            progress(0.5 * step / total_steps, desc="Restoring faces...")
+        enhanced = restore_faces(enhanced)
+        if progress:
+            step += 1
+            progress(0.5 * step / total_steps, desc="Upscaling...")
+        enhanced = advanced_tile_upscale(enhanced)
+        return enhanced
+    
+    if mode in (ENHANCE_MODE_CLEAN, ENHANCE_MODE_FULL_ENHANCE):
+        if progress:
+            step += 1
+            progress(0.7 * step / total_steps, desc="Removing artifacts...")
+        enhanced = remove_artifacts(enhanced)
+        if progress:
+            step += 1
+            progress(0.7 * step / total_steps, desc="Repairing skin and faces...")
+        enhanced = enhanced_skin_repair(enhanced)
+        enhanced = restore_faces(enhanced)
+    
+    if mode in (ENHANCE_MODE_UPSCALE, ENHANCE_MODE_CLEAN, ENHANCE_MODE_MAX_DETAIL, ENHANCE_MODE_FULL_ENHANCE):
+        if progress:
+            step += 1
+            progress(0.8 * step / total_steps, desc="Upscaling image...")
+        enhanced = advanced_tile_upscale(enhanced)
+    
+    if mode in (ENHANCE_MODE_MAX_DETAIL, ENHANCE_MODE_FULL_ENHANCE):
+        if progress:
+            step += 1
+            progress(0.9 * step / total_steps, desc="Enhancing details...")
+        enhanced = add_ultra_detail(enhanced, strength=0.7)
+        enhanced = apply_high_frequency_details(enhanced, amount=0.5)
+        enhanced = smart_sharpen(enhanced, strength=SMART_SHARPENING_STRENGTH)
+        if progress:
+            step += 1
+            progress(0.95 * step / total_steps, desc="Adding final grain...")
+        enhanced = add_film_grain(enhanced, seed)
+    
+    return enhanced
+
+# ============================================================================
+# UTILITY FUNCTIONS
+# ============================================================================
+
+def use_output_as_input(output_images):
+    """Move the first output image into the Image 1 slot."""
+    if not output_images:
+        return gr.update()
+    first = output_images[0]
+    path = first[0] if isinstance(first, (list, tuple)) else first
+    return gr.update(value=path)
+
+def check_gpu_memory():
+    """Check available GPU memory"""
+    if device == "cuda":
+        try:
+            total = torch.cuda.get_device_properties(0).total_memory
+            reserved = torch.cuda.memory_reserved(0)
+            allocated = torch.cuda.memory_allocated(0)
+            free = total - reserved
+            print(f"GPU Memory: Total={total/1024**3:.2f}GB, Reserved={reserved/1024**3:.2f}GB, Allocated={allocated/1024**3:.2f}GB, Free={free/1024**3:.2f}GB")
+            return free > 1024**3
+        except Exception as e:
+            print(f"Failed to check GPU memory: {e}")
+            return True
+    return True
+
+def clear_gpu_cache():
+    """Clear GPU cache to free up memory"""
+    if device == "cuda":
+        try:
+            torch.cuda.empty_cache()
+            import gc
+            gc.collect()
+            print("GPU cache cleared")
+        except Exception as e:
+            print(f"Failed to clear GPU cache: {e}")
+
+# ============================================================================
+# MAIN INFERENCE FUNCTION (Enhanced)
+# ============================================================================
+
+MAX_SEED = np.iinfo(np.int32).max
+
+@spaces.GPU(duration=60)
+def infer(
+    image_1,
+    image_2,
+    prompt,
+    seed=42,
+    randomize_seed=False,
+    true_guidance_scale=1.0,
+    num_inference_steps=4,
+    height=None,
+    width=None,
+    enhance_mode=ENHANCE_MODE_OFF,
+    num_images_per_prompt=1,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """
+    Enhanced image generation with advanced editing and enhancement options
+    """
+    negative_prompt = " "
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    pil_images = []
+    for img in (image_1, image_2):
+        if img is None:
+            continue
+        try:
+            if isinstance(img, str):
+                pil_images.append(Image.open(img).convert("RGB"))
+            elif isinstance(img, Image.Image):
+                pil_images.append(img.convert("RGB"))
+            elif hasattr(img, "name"):
+                pil_images.append(Image.open(img.name).convert("RGB"))
+        except Exception:
+            continue
+
+    if height==256 and width==256:
+        height, width = None, None
+    
+    print(f"Generation Parameters:")
+    print(f"   Prompt: '{prompt}'")
+    print(f"   Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}")
+    print(f"   Size: {width}x{height}, Enhance Mode: {enhance_mode}")
+
+    if not check_gpu_memory():
+        clear_gpu_cache()
+        if not check_gpu_memory():
+            raise gr.Error("Insufficient GPU memory. Please reduce image size or close other applications.")
+
+    try:
+        images_pil = pipe(
+            image=pil_images if len(pil_images) > 0 else None,
+            prompt=prompt,
+            height=height,
+            width=width,
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            true_cfg_scale=true_guidance_scale,
+            num_images_per_prompt=num_images_per_prompt,
+        ).images
+    except Exception as e:
+        clear_gpu_cache()
+        raise gr.Error(f"Image generation failed: {e}")
+
+    if enhance_mode != ENHANCE_MODE_OFF:
+        images_pil = [
+            apply_enhancement(img, enhance_mode, seed=seed + idx, progress=progress)
+            for idx, img in enumerate(images_pil)
+        ]
+
+    output_paths = []
+    os.makedirs("outputs", exist_ok=True)
+    for idx, img in enumerate(images_pil):
+        output_path = f"outputs/output_{seed}_{idx}_{int(time.time()*1000)}.png"
+        img.save(output_path)
+        output_paths.append(output_path)
+
+    clear_gpu_cache()
+
+    return output_paths, seed, gr.update(visible=True), gr.update(visible=bool(VIDEO_SPACE_ID))
+
+
+# ============================================================================
+# UI LAYOUT (Enhanced)
+# ============================================================================
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 1024px;
+}
+#logo-title {
+    text-align: center;
+}
+#logo-title h1 {
+    margin-bottom: 0;
+}
+#logo-title h2 {
+    color: #5b47d1;
+    font-style: italic;
+    margin-top: 0;
+}
+#edit_text{margin-top: -62px !important}
+.enhance-info {
+    font-size: 0.9em;
+    color: #666;
+    margin-top: 5px;
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.HTML(f"""
+        <div id="logo-title">
+            <h1>Pro Realism Edit Studio - Enhanced</h1>
+            <h2>Rapid Edit with Real-ESRGAN & Face Restoration</h2>
+        </div>
+        """)
+        
+        gr.Markdown("""
+        Powered by:
+        - [Qwen-Image-Edit-2511](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) 
+        - [Phr00t's Rapid-AIO v23](https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO) accelerated transformer
+        - [Real-ESRGAN](https://huggingface.co/ai-forever/Real-ESRGAN) for high-quality upscaling
+        - [GFPGAN](https://github.com/TencentARC/GFPGAN) for face restoration
+        
+        Upload an image and enter your prompt to edit it.
+        
+        Pro Tips:
+        - Use Face Enhance mode for portrait photography
+        - Use Max Detail for product shots and textures
+        - Use Full Enhance for comprehensive improvement
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    image_1 = gr.Image(label="Image 1", type="filepath", interactive=True)
+                    image_2 = gr.Image(label="Image 2 (optional)", type="filepath", interactive=True)
+
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=True,
+                    placeholder="Enter your prompt here...",
+                )
+                
+                enhance_mode = gr.Radio(
+                    label="Enhance Mode",
+                    choices=ENHANCE_MODE_CHOICES,
+                    value=ENHANCE_MODE_OFF,
+                    interactive=True,
+                    info="Choose enhancement level for your output"
+                )
+                
+                enhance_info = gr.Markdown("""
+                Enhancement Options:
+                - Off: No post-processing
+                - Upscale Only: 4x upscaling with Real-ESRGAN
+                - Clean & Restore: Artifact removal + skin/face restoration
+                - Max Detail: Full detail enhancement with sharpening
+                - Face Enhance: Specialized face restoration + upscaling
+                - Full Enhance: Complete pipeline (clean + detail + face + upscale)
+                """, visible=False)
+                
+                run_button = gr.Button("Generate!", variant="primary")
+                
+                with gr.Accordion("Advanced Settings", open=False):
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        step=1,
+                        value=0,
+                    )
+
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+                    with gr.Row():
+                        true_guidance_scale = gr.Slider(
+                            label="True guidance scale",
+                            minimum=1.0,
+                            maximum=10.0,
+                            step=0.1,
+                            value=1.0
+                        )
+
+                        num_inference_steps = gr.Slider(
+                            label="Number of inference steps",
+                            minimum=1,
+                            maximum=40,
+                            step=1,
+                            value=4,
+                        )
+                    
+                    with gr.Row():
+                        height = gr.Slider(
+                            label="Height",
+                            minimum=256,
+                            maximum=2048,
+                            step=8,
+                            value=None,
+                        )
+                        
+                        width = gr.Slider(
+                            label="Width",
+                            minimum=256,
+                            maximum=2048,
+                            step=8,
+                            value=None,
+                        )
+                    
+                    gr.Markdown("""
+                    Performance Tips:
+                    - Use 4 steps for fastest results
+                    - Increase steps (8-20) for better quality
+                    - Lower guidance scale for more creative freedom
+                    """)
+
+            with gr.Column():
+                result = gr.Gallery(label="Result", show_label=False, type="filepath")
+                with gr.Row():
+                    use_output_btn = gr.Button("Use as input", variant="secondary", size="sm", visible=False)
+                    turn_video_btn = gr.Button("Turn into Video", variant="secondary", size="sm", visible=False)
+                output_video = gr.Video(label="Generated Video", autoplay=True, visible=False)
+
+                with gr.Row():
+                    gr.Markdown("### History")
+                    clear_history_button = gr.Button("Clear History", size="sm", variant="stop")
+                
+                history_gallery = gr.Gallery(
+                    label="Click any image to use as input", 
+                    interactive=False,
+                    show_label=True,
+                    visible=True
+                )
+
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer,
+        inputs=[
+            image_1,
+            image_2,
+            prompt,
+            seed,
+            randomize_seed,
+            true_guidance_scale,
+            num_inference_steps,
+            height,
+            width,
+            enhance_mode,
+        ],
+        outputs=[result, seed, use_output_btn, turn_video_btn],
+
+    ).then(
+        fn=update_history,
+        inputs=[result, history_gallery],
+        outputs=history_gallery,
+    )
+
+    enhance_mode.change(
+        fn=lambda mode: gr.update(visible=mode != ENHANCE_MODE_OFF),
+        inputs=[enhance_mode],
+        outputs=[enhance_info]
+    )
+
+    use_output_btn.click(
+        fn=use_output_as_input,
+        inputs=[result],
+        outputs=[image_1]
+    )
+
+    history_gallery.select(
+        fn=use_history_as_input,
+        inputs=None,
+        outputs=[image_1],
+    )
+    
+    clear_history_button.click(
+        fn=lambda: [],
+        inputs=None,
+        outputs=history_gallery,
+    )
+
+    turn_video_btn.click(
+        fn=lambda: gr.update(visible=True),   
+        inputs=None,
+        outputs=[output_video],
+    ).then(
+        fn=turn_into_video,
+        inputs=[image_1, result, prompt],
+        outputs=[output_video],
+    )
+
+
+if __name__ == "__main__":
+    print(f"Device: {device}")
+    if device == "cuda":
+        print(f"GPU: {torch.cuda.get_device_name(0)}")
+    check_gpu_memory()
+    print(f"Starting Pro Realism Edit Studio v{APP_VERSION}")
+    demo.launch()

app_clean.py

@@ -0,0 +1,986 @@
+#!/usr/bin/env python3
+"""
+Pro Realism Edit Studio - Enhanced Edition
+=========================================
+
+Advanced image editing and enhancement studio powered by:
+- Qwen-Image-Edit-2511 with Phr00t's Rapid-AIO v23 accelerated transformer
+- Real-ESRGAN for high-quality upscaling
+- GFPGAN/CodeFormer for face restoration
+- Multi-stage detail enhancement pipeline
+
+Author: Enhanced with Hugging Face CLI and image generation expertise
+Version: 1.0.0
+"""
+
+import gradio as gr
+import numpy as np
+import random
+import torch
+import spaces
+import os
+import time
+import tempfile
+from pathlib import Path
+
+# Advanced imports
+from accelerate import init_empty_weights
+from collections import OrderedDict
+from PIL import Image, ImageEnhance, ImageFilter, ImageOps
+from diffusers.models import QwenImageTransformer2DModel as DiffusersQwenImageTransformer2DModel
+from diffusers.models.model_loading_utils import load_model_dict_into_meta
+from huggingface_hub import hf_hub_download, HfApi, login, whoami
+from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
+from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
+from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
+from safetensors import safe_open
+
+from gradio_client import Client, handle_file
+
+# ============================================================================
+# CONFIGURATION - Model IDs and Parameters
+# ============================================================================
+
+# Base model configuration
+BASE_MODEL_ID = "Qwen/Qwen-Image-Edit-2511"
+APP_VERSION = "1.0.0"
+PHR00T_REPO_ID = os.environ.get("PHR00T_REPO_ID", "Phr00t/Qwen-Image-Edit-Rapid-AIO").strip()
+RAPID_TRANSFORMER_FILENAME = os.environ.get(
+    "RAPID_TRANSFORMER_FILENAME",
+    "v23/Qwen-Rapid-AIO-NSFW-v23.safetensors",
+).strip()
+PHR00T_TRANSFORMER_PREFIX = "model.diffusion_model."
+VIDEO_SPACE_ID = os.environ.get("VIDEO_SPACE_ID", "").strip()
+
+# Enhanced Upscaler Configuration
+UPSCALER_MODEL_ID = os.environ.get("UPSCALER_MODEL_ID", "ai-forever/Real-ESRGAN").strip()
+UPSCALER_MODEL_FILENAME = os.environ.get("UPSCALER_MODEL_FILENAME", "RealESRGAN_x4plus.pth").strip()
+UPSCALER_TILE_SIZE = int(os.environ.get("UPSCALER_TILE_SIZE", "512"))
+UPSCALER_TILE_OVERLAP = int(os.environ.get("UPSCALER_TILE_OVERLAP", "64"))
+ENHANCE_MAX_INPUT_EDGE = int(os.environ.get("ENHANCE_MAX_INPUT_EDGE", "2048"))
+ENHANCE_GRAIN_STRENGTH = float(os.environ.get("ENHANCE_GRAIN_STRENGTH", "0.015"))
+
+# Face Restoration Configuration
+FACE_RESTORATION_MODEL = os.environ.get("FACE_RESTORATION_MODEL", "Xintao/GFPGAN").strip()
+FACE_RESTORATION_WEIGHTS = os.environ.get("FACE_RESTORATION_WEIGHTS", "GFPGANv1.3.pth").strip()
+
+# Advanced Detail Enhancement Configuration
+DETAIL_ENHANCEMENT_ENABLED = os.environ.get("DETAIL_ENHANCEMENT_ENABLED", "true").lower() == "true"
+SMART_SHARPENING_STRENGTH = float(os.environ.get("SMART_SHARPENING_STRENGTH", "1.15"))
+
+# ============================================================================
+# ENHANCEMENT MODES
+# ============================================================================
+
+ENHANCE_MODE_OFF = "Off"
+ENHANCE_MODE_UPSCALE = "Upscale Only"
+ENHANCE_MODE_CLEAN = "Clean & Restore"
+ENHANCE_MODE_MAX_DETAIL = "Max Detail"
+ENHANCE_MODE_FACE_ENHANCE = "Face Enhance"
+ENHANCE_MODE_FULL_ENHANCE = "Full Enhance"
+ENHANCE_MODE_CHOICES = [
+    ENHANCE_MODE_OFF, 
+    ENHANCE_MODE_UPSCALE, 
+    ENHANCE_MODE_CLEAN, 
+    ENHANCE_MODE_MAX_DETAIL,
+    ENHANCE_MODE_FACE_ENHANCE,
+    ENHANCE_MODE_FULL_ENHANCE
+]
+
+# ============================================================================
+# GLOBAL MODEL CACHE
+# ============================================================================
+
+_upscaler_model = None
+_face_restoration_model = None
+_detail_enhancement_model = None
+
+# ============================================================================
+# HUGGING FACE CLI EXPERT FUNCTIONS
+# ============================================================================
+
+def download_model_with_retry(repo_id, filename, max_retries=3):
+    """Download model with retry logic and error handling"""
+    for attempt in range(max_retries):
+        try:
+            return hf_hub_download(repo_id=repo_id, filename=filename)
+        except Exception as e:
+            if attempt == max_retries - 1:
+                raise RuntimeError(f"Failed to download {filename} from {repo_id} after {max_retries} attempts: {e}")
+            time.sleep(2 ** attempt)
+    return None
+
+# ============================================================================
+# VIDEO GENERATION (Preserved from original)
+# ============================================================================
+
+def turn_into_video(input_image, output_images, prompt, progress=gr.Progress(track_tqdm=True)):
+    """Convert image edit into video transition"""
+    if not VIDEO_SPACE_ID:
+        raise gr.Error("Video generation is not configured for this Space.")
+    if not input_image or not output_images:
+        raise gr.Error("Please generate an output image first.")
+
+    progress(0.02, desc="Preparing images...")
+
+    def extract_pil(img_entry):
+        if isinstance(img_entry, tuple) and isinstance(img_entry[0], Image.Image):
+            return img_entry[0]
+        elif isinstance(img_entry, Image.Image):
+            return img_entry
+        elif isinstance(img_entry, str):
+            return Image.open(img_entry)
+        else:
+            raise gr.Error(f"Unsupported image format: {type(img_entry)}")
+
+    start_img = extract_pil(input_image)
+    end_img = extract_pil(output_images[0])
+
+    progress(0.10, desc="Saving temp files...")
+
+    with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_start, \
+         tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp_end:
+        start_img.save(tmp_start.name)
+        end_img.save(tmp_end.name)
+
+    progress(0.20, desc="Connecting to video Space...")
+
+    client = Client(VIDEO_SPACE_ID)
+
+    progress(0.35, desc="Generating video...")
+
+    video_path, seed = client.predict(
+        start_image_pil=handle_file(tmp_start.name),
+        end_image_pil=handle_file(tmp_end.name),
+        prompt=prompt or "smooth cinematic transition",
+        api_name="/generate_video"
+    )
+
+    progress(0.95, desc="Finalizing...")
+    return video_path['video']
+
+
+# ============================================================================
+# HISTORY MANAGEMENT (Enhanced)
+# ============================================================================
+
+def update_history(new_images, history):
+    """Updates the history gallery with the new images."""
+    time.sleep(0.3)
+    if history is None:
+        history = []
+    if new_images is not None and len(new_images) > 0:
+        if not isinstance(history, list):
+            history = list(history) if history else []
+        for img in new_images:
+            history.insert(0, img)
+    history = history[:50]
+    return history
+
+def use_history_as_input(evt: gr.SelectData):
+    """Sets the selected history image into the Image 1 slot."""
+    if evt.value is not None:
+        return gr.update(value=evt.value)
+    return gr.update()
+
+# ============================================================================
+# MODEL LOADING (Enhanced with better error handling)
+# ============================================================================
+
+dtype = torch.bfloat16
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def load_phr00t_rapid_transformer(torch_dtype):
+    """Load Phr00t's Rapid-AIO v23 transformer with enhanced error handling"""
+    checkpoint_path = download_model_with_retry(PHR00T_REPO_ID, RAPID_TRANSFORMER_FILENAME)
+    
+    try:
+        config = DiffusersQwenImageTransformer2DModel.load_config(
+            BASE_MODEL_ID,
+            subfolder="transformer",
+        )
+    except Exception as e:
+        raise RuntimeError(f"Failed to load config for {BASE_MODEL_ID}: {e}")
+    
+    with init_empty_weights():
+        transformer = DiffusersQwenImageTransformer2DModel.from_config(config)
+
+    expected_keys = set(transformer.state_dict().keys())
+    state_dict = OrderedDict()
+    
+    try:
+        with safe_open(checkpoint_path, framework="pt", device="cpu") as checkpoint:
+            for key in checkpoint.keys():
+                if not key.startswith(PHR00T_TRANSFORMER_PREFIX):
+                    continue
+                mapped_key = key.removeprefix(PHR00T_TRANSFORMER_PREFIX)
+                if mapped_key in expected_keys:
+                    state_dict[mapped_key] = checkpoint.get_tensor(key)
+    except Exception as e:
+        raise RuntimeError(f"Failed to load checkpoint from {checkpoint_path}: {e}")
+
+    missing_keys = sorted(expected_keys.difference(state_dict.keys()))
+    if missing_keys:
+        sample = ", ".join(missing_keys[:20])
+        raise RuntimeError(
+            f"Phr00t Rapid-AIO transformer checkpoint is missing {len(missing_keys)} "
+            f"required diffusers keys after prefix conversion. First missing keys: {sample}"
+        )
+
+    try:
+        load_model_dict_into_meta(transformer, state_dict, dtype=torch_dtype)
+    except Exception as e:
+        raise RuntimeError(f"Failed to load state dict into meta: {e}")
+    
+    meta_parameters = [name for name, parameter in transformer.named_parameters() if parameter.is_meta]
+    if meta_parameters:
+        sample = ", ".join(meta_parameters[:20])
+        raise RuntimeError(
+            f"Phr00t Rapid-AIO transformer still has {len(meta_parameters)} meta parameters "
+            f"after loading. First meta parameters: {sample}"
+        )
+
+    transformer.eval()
+    return transformer
+
+# Load main pipeline
+try:
+    pipe = QwenImageEditPlusPipeline.from_pretrained(
+        BASE_MODEL_ID,
+        transformer=load_phr00t_rapid_transformer(dtype),
+        torch_dtype=dtype
+    ).to(device)
+    print("Successfully loaded Qwen-Image-Edit-2511 with Rapid-AIO v23 transformer")
+except Exception as e:
+    print(f"Failed to load main pipeline: {e}")
+    raise
+
+# Apply optimizations
+pipe.transformer.__class__ = QwenImageTransformer2DModel
+pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+print("Applied FA3 attention processor optimization")
+
+# ============================================================================
+# ENHANCED UPSCALER (Real-ESRGAN based)
+# ============================================================================
+
+def load_upscaler_model():
+    """Load Real-ESRGAN model for high-quality upscaling"""
+    global _upscaler_model
+    if _upscaler_model is not None:
+        return _upscaler_model
+
+    try:
+        import spandrel
+        import spandrel_extra_arches
+        spandrel_extra_arches.install()
+    except ImportError as exc:
+        raise gr.Error("Enhance mode requires spandrel and spandrel_extra_arches to be installed. "
+                       "Install with: pip install spandrel spandrel_extra_arches") from exc
+
+    try:
+        model_path = download_model_with_retry(UPSCALER_MODEL_ID, UPSCALER_MODEL_FILENAME)
+        model = spandrel.ModelLoader().load_from_file(model_path)
+        model.eval().to(device)
+        _upscaler_model = model
+        print(f"Successfully loaded upscaler: {UPSCALER_MODEL_ID}/{UPSCALER_MODEL_FILENAME}")
+        return _upscaler_model
+    except Exception as e:
+        print(f"Failed to load upscaler model: {e}")
+        print("Falling back to Nomos upscaler...")
+        try:
+            model_path = download_model_with_retry("Phips/4xNomos8k_atd_jpg", "4xNomos8k_atd_jpg.safetensors")
+            model = spandrel.ModelLoader().load_from_file(model_path)
+            model.eval().to(device)
+            _upscaler_model = model
+            return _upscaler_model
+        except Exception as fallback_error:
+            raise gr.Error(f"Failed to load all upscaler models: {e} | {fallback_error}")
+
+def image_to_tensor(image):
+    """Convert PIL Image to tensor"""
+    array = np.asarray(image.convert("RGB")).astype(np.float32) / 255.0
+    tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0)
+    return tensor.to(device)
+
+def tensor_to_image(tensor):
+    """Convert tensor to PIL Image"""
+    array = tensor.squeeze(0).detach().float().cpu().clamp(0, 1).permute(1, 2, 0).numpy()
+    return Image.fromarray((array * 255.0).round().astype(np.uint8), mode="RGB")
+
+def validate_enhance_input_size(image):
+    """Validate image size for enhancement"""
+    max_edge = max(image.size)
+    if max_edge > ENHANCE_MAX_INPUT_EDGE:
+        raise gr.Error(
+            f"Enhance mode accepts images up to {ENHANCE_MAX_INPUT_EDGE}px on the longest edge. "
+            f"Current image is {image.width}x{image.height}. "
+            f"Consider resizing your image first."
+        )
+
+def advanced_tile_upscale(image, scale=4):
+    """
+    Advanced tiling upscaler with improved blending and edge handling
+    Uses Real-ESRGAN for superior quality compared to Nomos
+    """
+    validate_enhance_input_size(image)
+    model = load_upscaler_model()
+    tensor = image_to_tensor(image)
+    _, _, height, width = tensor.shape
+    
+    base_tile_size = UPSCALER_TILE_SIZE
+    optimal_tile_size = min(base_tile_size, max(height, width) // 2)
+    tile_size = max(64, optimal_tile_size)
+    overlap = max(0, min(UPSCALER_TILE_OVERLAP, tile_size // 2))
+    step = max(1, tile_size - overlap)
+    
+    y_positions = list(range(0, height, step))
+    if y_positions[-1] + tile_size < height:
+        y_positions.append(max(0, height - tile_size))
+    
+    x_positions = list(range(0, width, step))
+    if x_positions[-1] + tile_size < width:
+        x_positions.append(max(0, width - tile_size))
+    
+    y_positions = sorted(set(y_positions))
+    x_positions = sorted(set(x_positions))
+    
+    output = None
+    weights = None
+
+    with torch.inference_mode():
+        for y in y_positions:
+            for x in x_positions:
+                y1 = min(y + tile_size, height)
+                x1 = min(x + tile_size, width)
+                tile = tensor[:, :, y:y1, x:x1]
+                
+                upscaled_tile = model(tile).clamp(0, 1)
+                
+                scale_y = upscaled_tile.shape[-2] // tile.shape[-2]
+                scale_x = upscaled_tile.shape[-1] // tile.shape[-1]
+                
+                if output is None:
+                    output = torch.zeros(
+                        (1, 3, height * scale_y, width * scale_x),
+                        dtype=upscaled_tile.dtype,
+                        device=upscaled_tile.device,
+                    )
+                    weights = torch.zeros_like(output)
+                
+                oy0, oy1 = y * scale_y, y1 * scale_y
+                ox0, ox1 = x * scale_x, x1 * scale_x
+                output[:, :, oy0:oy1, ox0:ox1] += upscaled_tile
+                weights[:, :, oy0:oy1, ox0:ox1] += 1
+
+    output = output / weights.clamp_min(1)
+    return tensor_to_image(output)
+
+# ============================================================================
+# ENHANCED DETAILER (Multi-stage processing)
+# ============================================================================
+
+def smart_sharpen(image, strength=1.15):
+    """
+    Smart sharpening with edge detection to avoid oversharpening smooth areas
+    """
+    if strength <= 0:
+        return image
+    
+    img_array = np.array(image.convert("RGB"))
+    
+    if strength > 1.0:
+        enhanced = ImageEnhance.Sharpness(image).enhance(strength)
+        
+        gray = image.convert("L")
+        edges = gray.filter(ImageFilter.FIND_EDGES)
+        edge_mask = edges.filter(ImageFilter.GaussianBlur(radius=1))
+        edge_mask = edge_mask.point(lambda x: min(x * 0.3, 255))
+        
+        sharpened_array = np.array(enhanced)
+        original_array = img_array
+        edge_array = np.array(edge_mask).astype(float) / 255.0
+        
+        for c in range(3):
+            sharpened_array[:, :, c] = (
+                edge_array * sharpened_array[:, :, c] + 
+                (1 - edge_array) * original_array[:, :, c]
+            )
+        
+        image = Image.fromarray(np.clip(sharpened_array, 0, 255).astype(np.uint8))
+    
+    return image
+
+def add_ultra_detail(image, strength=0.8):
+    """
+    Add ultra-fine details using high-frequency enhancement
+    """
+    if strength <= 0:
+        return image
+    
+    original = image.convert("RGB")
+    blurred = original.filter(ImageFilter.GaussianBlur(radius=2))
+    
+    high_freq = ImageChops.subtract(original, blurred)
+    high_freq_enhanced = ImageEnhance.Contrast(high_freq).enhance(1.0 + strength)
+    result = ImageChops.add(original, high_freq_enhanced)
+    
+    return result
+
+def apply_high_frequency_details(image, amount=0.6):
+    """
+    Apply high-frequency detail enhancement for crisp textures
+    """
+    if amount <= 0:
+        return image
+    
+    scales = [1, 2, 4]
+    result = image.convert("RGB")
+    
+    for scale in scales:
+        blurred = result.filter(ImageFilter.GaussianBlur(radius=scale))
+        high_freq = ImageChops.subtract(result, blurred)
+        enhanced_hf = ImageEnhance.Contrast(high_freq).enhance(1.0 + amount * 0.3)
+        result = ImageChops.add(result, enhanced_hf)
+    
+    return result
+
+# ============================================================================
+# ENHANCED CLEANER (Face Restoration + Artifact Removal)
+# ============================================================================
+
+def load_face_restoration_model():
+    """Load GFPGAN model for face restoration"""
+    global _face_restoration_model
+    if _face_restoration_model is not None:
+        return _face_restoration_model
+    
+    try:
+        from gfpgan import GFPGANer
+        model_path = download_model_with_retry(FACE_RESTORATION_MODEL, FACE_RESTORATION_WEIGHTS)
+        
+        restorer = GFPGANer(
+            model_path=model_path,
+            upscale=1,
+            arch='clean',
+            channel_multiplier=2,
+            bg_upsampler=None
+        )
+        
+        _face_restoration_model = restorer
+        print("Successfully loaded GFPGAN face restoration model")
+        return _face_restoration_model
+        
+    except ImportError:
+        print("GFPGAN not available, face restoration will use fallback methods")
+        return None
+    except Exception as e:
+        print(f"Failed to load face restoration model: {e}")
+        return None
+
+def restore_faces(image):
+    """Restore faces in an image using GFPGAN"""
+    restorer = load_face_restoration_model()
+    if restorer is None:
+        print("Face restoration model not available, using skin repair fallback")
+        return repair_skin_texture(image)
+    
+    try:
+        img_array = np.array(image.convert("RGB"))
+        restored_array, _ = restorer.enhance(img_array, has_aligned=False, only_center_face=False, paste_back=True)
+        restored_image = Image.fromarray(restored_array.astype(np.uint8))
+        return restored_image
+    except Exception as e:
+        print(f"Face restoration failed: {e}, using skin repair fallback")
+        return repair_skin_texture(image)
+
+def remove_artifacts(image):
+    """Remove compression artifacts and noise"""
+    denoised = image.filter(ImageFilter.MedianFilter(size=3))
+    smoothed = denoised.filter(ImageFilter.GaussianBlur(radius=0.5))
+    result = Image.blend(image, smoothed, alpha=0.3)
+    return result
+
+def enhanced_skin_repair(image):
+    """Enhanced skin repair with better color detection and blending"""
+    base = image.convert("RGB")
+    
+    ycbcr = np.asarray(base.convert("YCbCr"))
+    y, cb, cr = ycbcr[:, :, 0], ycbcr[:, :, 1], ycbcr[:, :, 2]
+    
+    skin_mask = (
+        (cr > 130) & (cr < 170) &
+        (cb > 70) & (cb < 140) &
+        (y > 80)
+    ).astype(np.uint8) * 255
+    
+    try:
+        import cv2
+        kernel = np.ones((5, 5), np.uint8)
+        skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_OPEN, kernel)
+        skin_mask = cv2.morphologyEx(skin_mask, cv2.MORPH_CLOSE, kernel)
+        skin_mask = cv2.GaussianBlur(skin_mask, (7, 7), 0)
+    except ImportError:
+        from scipy import ndimage
+        skin_mask = ndimage.binary_opening(skin_mask > 128, structure=np.ones((3, 3))).astype(np.uint8) * 255
+        skin_mask = ndimage.gaussian_filter(skin_mask, sigma=3)
+    
+    mask_image = Image.fromarray(skin_mask, mode="L")
+    repaired = base.filter(ImageFilter.MedianFilter(size=3))
+    repaired = repaired.filter(ImageFilter.GaussianBlur(radius=0.4))
+    non_skin = ImageOps.invert(mask_image)
+    sharpened = ImageEnhance.Sharpness(base).enhance(1.15)
+    blended = Image.composite(repaired, sharpened, mask_image)
+    result = ImageEnhance.Sharpness(blended).enhance(1.05)
+    return result
+
+# Original skin repair functions (preserved for compatibility)
+def skin_repair_mask(image):
+    ycbcr = np.asarray(image.convert("YCbCr"))
+    cb = ycbcr[:, :, 1]
+    cr = ycbcr[:, :, 2]
+    mask = (
+        (cr >= 135)
+        & (cr <= 180)
+        & (cb >= 75)
+        & (cb <= 135)
+    ).astype(np.uint8) * 255
+    mask_image = Image.fromarray(mask, mode="L")
+    return mask_image.filter(ImageFilter.GaussianBlur(radius=1.2))
+
+def repair_skin_texture(image):
+    base = image.convert("RGB")
+    mask = skin_repair_mask(base)
+    repaired = base.filter(ImageFilter.MedianFilter(size=3)).filter(ImageFilter.GaussianBlur(radius=0.35))
+    blended = Image.composite(repaired, base, mask)
+    return ImageEnhance.Sharpness(blended).enhance(1.08)
+
+def add_film_grain(image, seed):
+    base = image.convert("RGB")
+    array = np.asarray(base).astype(np.float32)
+    rng = np.random.default_rng(seed)
+    grain = rng.normal(0.0, 255.0 * ENHANCE_GRAIN_STRENGTH, size=(array.shape[0], array.shape[1], 1))
+    array = np.clip(array + grain, 0, 255)
+    return Image.fromarray(array.astype(np.uint8), mode="RGB")
+
+# ============================================================================
+# ENHANCED APPLY ENHANCEMENT (Main enhancement pipeline)
+# ============================================================================
+
+def apply_enhancement(image, enhance_mode, seed=0, progress=None):
+    """
+    Apply various enhancement modes to the image
+    """
+    mode = enhance_mode or ENHANCE_MODE_OFF
+    if mode not in ENHANCE_MODE_CHOICES:
+        raise gr.Error(f"Unknown enhance mode: {mode}")
+    if mode == ENHANCE_MODE_OFF:
+        return image
+
+    enhanced = image.convert("RGB")
+    
+    total_steps = 0
+    if mode == ENHANCE_MODE_UPSCALE:
+        total_steps = 1
+    elif mode == ENHANCE_MODE_CLEAN:
+        total_steps = 3
+    elif mode == ENHANCE_MODE_MAX_DETAIL:
+        total_steps = 4
+    elif mode == ENHANCE_MODE_FACE_ENHANCE:
+        total_steps = 2
+    elif mode == ENHANCE_MODE_FULL_ENHANCE:
+        total_steps = 5
+    
+    step = 0
+    
+    if mode == ENHANCE_MODE_FACE_ENHANCE:
+        if progress:
+            step += 1
+            progress(0.5 * step / total_steps, desc="Restoring faces...")
+        enhanced = restore_faces(enhanced)
+        if progress:
+            step += 1
+            progress(0.5 * step / total_steps, desc="Upscaling...")
+        enhanced = advanced_tile_upscale(enhanced)
+        return enhanced
+    
+    if mode in (ENHANCE_MODE_CLEAN, ENHANCE_MODE_FULL_ENHANCE):
+        if progress:
+            step += 1
+            progress(0.7 * step / total_steps, desc="Removing artifacts...")
+        enhanced = remove_artifacts(enhanced)
+        if progress:
+            step += 1
+            progress(0.7 * step / total_steps, desc="Repairing skin and faces...")
+        enhanced = enhanced_skin_repair(enhanced)
+        enhanced = restore_faces(enhanced)
+    
+    if mode in (ENHANCE_MODE_UPSCALE, ENHANCE_MODE_CLEAN, ENHANCE_MODE_MAX_DETAIL, ENHANCE_MODE_FULL_ENHANCE):
+        if progress:
+            step += 1
+            progress(0.8 * step / total_steps, desc="Upscaling image...")
+        enhanced = advanced_tile_upscale(enhanced)
+    
+    if mode in (ENHANCE_MODE_MAX_DETAIL, ENHANCE_MODE_FULL_ENHANCE):
+        if progress:
+            step += 1
+            progress(0.9 * step / total_steps, desc="Enhancing details...")
+        enhanced = add_ultra_detail(enhanced, strength=0.7)
+        enhanced = apply_high_frequency_details(enhanced, amount=0.5)
+        enhanced = smart_sharpen(enhanced, strength=SMART_SHARPENING_STRENGTH)
+        if progress:
+            step += 1
+            progress(0.95 * step / total_steps, desc="Adding final grain...")
+        enhanced = add_film_grain(enhanced, seed)
+    
+    return enhanced
+
+# ============================================================================
+# UTILITY FUNCTIONS
+# ============================================================================
+
+def use_output_as_input(output_images):
+    """Move the first output image into the Image 1 slot."""
+    if not output_images:
+        return gr.update()
+    first = output_images[0]
+    path = first[0] if isinstance(first, (list, tuple)) else first
+    return gr.update(value=path)
+
+def check_gpu_memory():
+    """Check available GPU memory"""
+    if device == "cuda":
+        try:
+            total = torch.cuda.get_device_properties(0).total_memory
+            reserved = torch.cuda.memory_reserved(0)
+            allocated = torch.cuda.memory_allocated(0)
+            free = total - reserved
+            print(f"GPU Memory: Total={total/1024**3:.2f}GB, Reserved={reserved/1024**3:.2f}GB, Allocated={allocated/1024**3:.2f}GB, Free={free/1024**3:.2f}GB")
+            return free > 1024**3
+        except Exception as e:
+            print(f"Failed to check GPU memory: {e}")
+            return True
+    return True
+
+def clear_gpu_cache():
+    """Clear GPU cache to free up memory"""
+    if device == "cuda":
+        try:
+            torch.cuda.empty_cache()
+            import gc
+            gc.collect()
+            print("GPU cache cleared")
+        except Exception as e:
+            print(f"Failed to clear GPU cache: {e}")
+
+# ============================================================================
+# MAIN INFERENCE FUNCTION (Enhanced)
+# ============================================================================
+
+MAX_SEED = np.iinfo(np.int32).max
+
+@spaces.GPU(duration=60)
+def infer(
+    image_1,
+    image_2,
+    prompt,
+    seed=42,
+    randomize_seed=False,
+    true_guidance_scale=1.0,
+    num_inference_steps=4,
+    height=None,
+    width=None,
+    enhance_mode=ENHANCE_MODE_OFF,
+    num_images_per_prompt=1,
+    progress=gr.Progress(track_tqdm=True),
+):
+    """
+    Enhanced image generation with advanced editing and enhancement options
+    """
+    negative_prompt = " "
+
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    pil_images = []
+    for img in (image_1, image_2):
+        if img is None:
+            continue
+        try:
+            if isinstance(img, str):
+                pil_images.append(Image.open(img).convert("RGB"))
+            elif isinstance(img, Image.Image):
+                pil_images.append(img.convert("RGB"))
+            elif hasattr(img, "name"):
+                pil_images.append(Image.open(img.name).convert("RGB"))
+        except Exception:
+            continue
+
+    if height==256 and width==256:
+        height, width = None, None
+    
+    print(f"Generation Parameters:")
+    print(f"   Prompt: '{prompt}'")
+    print(f"   Seed: {seed}, Steps: {num_inference_steps}, Guidance: {true_guidance_scale}")
+    print(f"   Size: {width}x{height}, Enhance Mode: {enhance_mode}")
+
+    if not check_gpu_memory():
+        clear_gpu_cache()
+        if not check_gpu_memory():
+            raise gr.Error("Insufficient GPU memory. Please reduce image size or close other applications.")
+
+    try:
+        images_pil = pipe(
+            image=pil_images if len(pil_images) > 0 else None,
+            prompt=prompt,
+            height=height,
+            width=width,
+            negative_prompt=negative_prompt,
+            num_inference_steps=num_inference_steps,
+            generator=generator,
+            true_cfg_scale=true_guidance_scale,
+            num_images_per_prompt=num_images_per_prompt,
+        ).images
+    except Exception as e:
+        clear_gpu_cache()
+        raise gr.Error(f"Image generation failed: {e}")
+
+    if enhance_mode != ENHANCE_MODE_OFF:
+        images_pil = [
+            apply_enhancement(img, enhance_mode, seed=seed + idx, progress=progress)
+            for idx, img in enumerate(images_pil)
+        ]
+
+    output_paths = []
+    os.makedirs("outputs", exist_ok=True)
+    for idx, img in enumerate(images_pil):
+        output_path = f"outputs/output_{seed}_{idx}_{int(time.time()*1000)}.png"
+        img.save(output_path)
+        output_paths.append(output_path)
+
+    clear_gpu_cache()
+
+    return output_paths, seed, gr.update(visible=True), gr.update(visible=bool(VIDEO_SPACE_ID))
+
+
+# ============================================================================
+# UI LAYOUT (Enhanced)
+# ============================================================================
+
+css = """
+#col-container {
+    margin: 0 auto;
+    max-width: 1024px;
+}
+#logo-title {
+    text-align: center;
+}
+#logo-title h1 {
+    margin-bottom: 0;
+}
+#logo-title h2 {
+    color: #5b47d1;
+    font-style: italic;
+    margin-top: 0;
+}
+#edit_text{margin-top: -62px !important}
+.enhance-info {
+    font-size: 0.9em;
+    color: #666;
+    margin-top: 5px;
+}
+"""
+
+with gr.Blocks(css=css) as demo:
+    with gr.Column(elem_id="col-container"):
+        gr.HTML(f"""
+        <div id="logo-title">
+            <h1>Pro Realism Edit Studio - Enhanced</h1>
+            <h2>Rapid Edit with Real-ESRGAN & Face Restoration</h2>
+        </div>
+        """)
+        
+        gr.Markdown("""
+        Powered by:
+        - [Qwen-Image-Edit-2511](https://huggingface.co/Qwen/Qwen-Image-Edit-2511) 
+        - [Phr00t's Rapid-AIO v23](https://huggingface.co/Phr00t/Qwen-Image-Edit-Rapid-AIO) accelerated transformer
+        - [Real-ESRGAN](https://huggingface.co/ai-forever/Real-ESRGAN) for high-quality upscaling
+        - [GFPGAN](https://github.com/TencentARC/GFPGAN) for face restoration
+        
+        Upload an image and enter your prompt to edit it.
+        
+        Pro Tips:
+        - Use Face Enhance mode for portrait photography
+        - Use Max Detail for product shots and textures
+        - Use Full Enhance for comprehensive improvement
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                with gr.Row():
+                    image_1 = gr.Image(label="Image 1", type="filepath", interactive=True)
+                    image_2 = gr.Image(label="Image 2 (optional)", type="filepath", interactive=True)
+
+                prompt = gr.Text(
+                    label="Prompt",
+                    show_label=True,
+                    placeholder="Enter your prompt here...",
+                )
+                
+                enhance_mode = gr.Radio(
+                    label="Enhance Mode",
+                    choices=ENHANCE_MODE_CHOICES,
+                    value=ENHANCE_MODE_OFF,
+                    interactive=True,
+                    info="Choose enhancement level for your output"
+                )
+                
+                enhance_info = gr.Markdown("""
+                Enhancement Options:
+                - Off: No post-processing
+                - Upscale Only: 4x upscaling with Real-ESRGAN
+                - Clean & Restore: Artifact removal + skin/face restoration
+                - Max Detail: Full detail enhancement with sharpening
+                - Face Enhance: Specialized face restoration + upscaling
+                - Full Enhance: Complete pipeline (clean + detail + face + upscale)
+                """, visible=False)
+                
+                run_button = gr.Button("Generate!", variant="primary")
+                
+                with gr.Accordion("Advanced Settings", open=False):
+                    seed = gr.Slider(
+                        label="Seed",
+                        minimum=0,
+                        maximum=MAX_SEED,
+                        step=1,
+                        value=0,
+                    )
+
+                    randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+
+                    with gr.Row():
+                        true_guidance_scale = gr.Slider(
+                            label="True guidance scale",
+                            minimum=1.0,
+                            maximum=10.0,
+                            step=0.1,
+                            value=1.0
+                        )
+
+                        num_inference_steps = gr.Slider(
+                            label="Number of inference steps",
+                            minimum=1,
+                            maximum=40,
+                            step=1,
+                            value=4,
+                        )
+                    
+                    with gr.Row():
+                        height = gr.Slider(
+                            label="Height",
+                            minimum=256,
+                            maximum=2048,
+                            step=8,
+                            value=None,
+                        )
+                        
+                        width = gr.Slider(
+                            label="Width",
+                            minimum=256,
+                            maximum=2048,
+                            step=8,
+                            value=None,
+                        )
+                    
+                    gr.Markdown("""
+                    Performance Tips:
+                    - Use 4 steps for fastest results
+                    - Increase steps (8-20) for better quality
+                    - Lower guidance scale for more creative freedom
+                    """)
+
+            with gr.Column():
+                result = gr.Gallery(label="Result", show_label=False, type="filepath")
+                with gr.Row():
+                    use_output_btn = gr.Button("Use as input", variant="secondary", size="sm", visible=False)
+                    turn_video_btn = gr.Button("Turn into Video", variant="secondary", size="sm", visible=False)
+                output_video = gr.Video(label="Generated Video", autoplay=True, visible=False)
+
+                with gr.Row():
+                    gr.Markdown("### History")
+                    clear_history_button = gr.Button("Clear History", size="sm", variant="stop")
+                
+                history_gallery = gr.Gallery(
+                    label="Click any image to use as input", 
+                    interactive=False,
+                    show_label=True,
+                    visible=True
+                )
+
+    gr.on(
+        triggers=[run_button.click, prompt.submit],
+        fn=infer,
+        inputs=[
+            image_1,
+            image_2,
+            prompt,
+            seed,
+            randomize_seed,
+            true_guidance_scale,
+            num_inference_steps,
+            height,
+            width,
+            enhance_mode,
+        ],
+        outputs=[result, seed, use_output_btn, turn_video_btn],
+
+    ).then(
+        fn=update_history,
+        inputs=[result, history_gallery],
+        outputs=history_gallery,
+    )
+
+    enhance_mode.change(
+        fn=lambda mode: gr.update(visible=mode != ENHANCE_MODE_OFF),
+        inputs=[enhance_mode],
+        outputs=[enhance_info]
+    )
+
+    use_output_btn.click(
+        fn=use_output_as_input,
+        inputs=[result],
+        outputs=[image_1]
+    )
+
+    history_gallery.select(
+        fn=use_history_as_input,
+        inputs=None,
+        outputs=[image_1],
+    )
+    
+    clear_history_button.click(
+        fn=lambda: [],
+        inputs=None,
+        outputs=history_gallery,
+    )
+
+    turn_video_btn.click(
+        fn=lambda: gr.update(visible=True),   
+        inputs=None,
+        outputs=[output_video],
+    ).then(
+        fn=turn_into_video,
+        inputs=[image_1, result, prompt],
+        outputs=[output_video],
+    )
+
+
+if __name__ == "__main__":
+    print(f"Device: {device}")
+    if device == "cuda":
+        print(f"GPU: {torch.cuda.get_device_name(0)}")
+    check_gpu_memory()
+    print(f"Starting Pro Realism Edit Studio v{APP_VERSION}")
+    demo.launch()

optimization.py

@@ -1,85 +1,85 @@
-"""
-AOT compilation optimization for Qwen-Image-Edit pipeline.
-"""
-
-import gc
-from typing import Any
-from typing import Callable
-from typing import ParamSpec
-from torchao.quantization import quantize_
-from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
-import spaces
-import torch
-from torch.utils._pytree import tree_map
-
-
-P = ParamSpec('P')
-
-
-TRANSFORMER_IMAGE_SEQ_LENGTH_DIM = torch.export.Dim('image_seq_length')
-TRANSFORMER_TEXT_SEQ_LENGTH_DIM = torch.export.Dim('text_seq_length')
-
-TRANSFORMER_DYNAMIC_SHAPES = {
-    'hidden_states': {
-        1: TRANSFORMER_IMAGE_SEQ_LENGTH_DIM,
-    },
-    'encoder_hidden_states': {
-        1: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
-    },
-    'encoder_hidden_states_mask': {
-        1: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
-    },
-    'image_rotary_emb': ({
-        0: TRANSFORMER_IMAGE_SEQ_LENGTH_DIM,
-    }, {
-        0: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
-    }),
-}
-
-
-INDUCTOR_CONFIGS = {
-    'conv_1x1_as_mm': True,
-    'epilogue_fusion': False,
-    'coordinate_descent_tuning': True,
-    'coordinate_descent_check_all_directions': True,
-    'max_autotune': True,
-    'triton.cudagraphs': True,
-}
-
-
-def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
-
-    @spaces.GPU(duration=1000)
-    def compile_transformer():
-
-        # Run warmup pass to capture transformer inputs (needs all components on GPU)
-        with spaces.aoti_capture(pipeline.transformer) as call:
-            pipeline(*args, **kwargs)
-
-        # Offload text encoder to CPU to free ~16GB during memory-intensive torch.export
-        # This happens AFTER warmup but BEFORE export to avoid device mismatch
-        text_encoder_device = next(pipeline.text_encoder.parameters()).device
-        pipeline.text_encoder.to('cpu')
-        gc.collect()
-        torch.cuda.empty_cache()
-
-        dynamic_shapes = tree_map(lambda t: None, call.kwargs)
-        dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
-
-        # quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
-
-        exported = torch.export.export(
-            mod=pipeline.transformer,
-            args=call.args,
-            kwargs=call.kwargs,
-            dynamic_shapes=dynamic_shapes,
-        )
-
-        compiled = spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
-
-        # Move text encoder back to original device
-        pipeline.text_encoder.to(text_encoder_device)
-
-        return compiled
-
+"""
+AOT compilation optimization for Qwen-Image-Edit pipeline.
+"""
+
+import gc
+from typing import Any
+from typing import Callable
+from typing import ParamSpec
+from torchao.quantization import quantize_
+from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
+import spaces
+import torch
+from torch.utils._pytree import tree_map
+
+
+P = ParamSpec('P')
+
+
+TRANSFORMER_IMAGE_SEQ_LENGTH_DIM = torch.export.Dim('image_seq_length')
+TRANSFORMER_TEXT_SEQ_LENGTH_DIM = torch.export.Dim('text_seq_length')
+
+TRANSFORMER_DYNAMIC_SHAPES = {
+    'hidden_states': {
+        1: TRANSFORMER_IMAGE_SEQ_LENGTH_DIM,
+    },
+    'encoder_hidden_states': {
+        1: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
+    },
+    'encoder_hidden_states_mask': {
+        1: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
+    },
+    'image_rotary_emb': ({
+        0: TRANSFORMER_IMAGE_SEQ_LENGTH_DIM,
+    }, {
+        0: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
+    }),
+}
+
+
+INDUCTOR_CONFIGS = {
+    'conv_1x1_as_mm': True,
+    'epilogue_fusion': False,
+    'coordinate_descent_tuning': True,
+    'coordinate_descent_check_all_directions': True,
+    'max_autotune': True,
+    'triton.cudagraphs': True,
+}
+
+
+def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
+
+    @spaces.GPU(duration=1000)
+    def compile_transformer():
+
+        # Run warmup pass to capture transformer inputs (needs all components on GPU)
+        with spaces.aoti_capture(pipeline.transformer) as call:
+            pipeline(*args, **kwargs)
+
+        # Offload text encoder to CPU to free ~16GB during memory-intensive torch.export
+        # This happens AFTER warmup but BEFORE export to avoid device mismatch
+        text_encoder_device = next(pipeline.text_encoder.parameters()).device
+        pipeline.text_encoder.to('cpu')
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        dynamic_shapes = tree_map(lambda t: None, call.kwargs)
+        dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
+
+        # quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
+
+        exported = torch.export.export(
+            mod=pipeline.transformer,
+            args=call.args,
+            kwargs=call.kwargs,
+            dynamic_shapes=dynamic_shapes,
+        )
+
+        compiled = spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
+
+        # Move text encoder back to original device
+        pipeline.text_encoder.to(text_encoder_device)
+
+        return compiled
+
     spaces.aoti_apply(compile_transformer(), pipeline.transformer)

qwenimage/pipeline_qwenimage_edit_plus.py

@@ -1,889 +1,889 @@
-# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import inspect
-import math
-from typing import Any, Callable, Dict, List, Optional, Union
-
-import numpy as np
-import torch
-from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor
-
-from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
-from diffusers.loaders import QwenImageLoraLoaderMixin
-from diffusers.models import AutoencoderKLQwenImage, QwenImageTransformer2DModel
-from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
-from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
-from diffusers.utils.torch_utils import randn_tensor
-from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from diffusers.pipelines.qwenimage.pipeline_output import QwenImagePipelineOutput
-
-
-if is_torch_xla_available():
-    import torch_xla.core.xla_model as xm
-
-    XLA_AVAILABLE = True
-else:
-    XLA_AVAILABLE = False
-
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-EXAMPLE_DOC_STRING = """
-    Examples:
-        ```py
-        >>> import torch
-        >>> from PIL import Image
-        >>> from diffusers import QwenImageEditPlusPipeline
-        >>> from diffusers.utils import load_image
-
-        >>> pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2511", torch_dtype=torch.bfloat16)
-        >>> pipe.to("cuda")
-        >>> image = load_image("input.png").convert("RGB")
-        >>> prompt = (
-        ...     "Make Pikachu hold a sign that says 'Qwen Edit is awesome', yarn art style, detailed, vibrant colors"
-        ... )
-        >>> # Depending on the variant being used, the pipeline call will slightly vary.
-        >>> # Refer to the pipeline documentation for more details.
-        >>> image = pipe(image, prompt, num_inference_steps=50).images[0]
-        >>> image.save("qwenimage_edit_plus.png")
-        ```
-"""
-
-CONDITION_IMAGE_SIZE = 384 * 384
-VAE_IMAGE_SIZE = 1024 * 1024
-
-
-# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.calculate_shift
-def calculate_shift(
-    image_seq_len,
-    base_seq_len: int = 256,
-    max_seq_len: int = 4096,
-    base_shift: float = 0.5,
-    max_shift: float = 1.15,
-):
-    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
-    b = base_shift - m * base_seq_len
-    mu = image_seq_len * m + b
-    return mu
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
-def retrieve_timesteps(
-    scheduler,
-    num_inference_steps: Optional[int] = None,
-    device: Optional[Union[str, torch.device]] = None,
-    timesteps: Optional[List[int]] = None,
-    sigmas: Optional[List[float]] = None,
-    **kwargs,
-):
-    r"""
-    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
-    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
-
-    Args:
-        scheduler (`SchedulerMixin`):
-            The scheduler to get timesteps from.
-        num_inference_steps (`int`):
-            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
-            must be `None`.
-        device (`str` or `torch.device`, *optional*):
-            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
-        timesteps (`List[int]`, *optional*):
-            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
-            `num_inference_steps` and `sigmas` must be `None`.
-        sigmas (`List[float]`, *optional*):
-            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
-            `num_inference_steps` and `timesteps` must be `None`.
-
-    Returns:
-        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
-        second element is the number of inference steps.
-    """
-    if timesteps is not None and sigmas is not None:
-        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
-    if timesteps is not None:
-        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accepts_timesteps:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" timestep schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    elif sigmas is not None:
-        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
-        if not accept_sigmas:
-            raise ValueError(
-                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
-                f" sigmas schedules. Please check whether you are using the correct scheduler."
-            )
-        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-        num_inference_steps = len(timesteps)
-    else:
-        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
-        timesteps = scheduler.timesteps
-    return timesteps, num_inference_steps
-
-
-# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
-def retrieve_latents(
-    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
-):
-    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
-        return encoder_output.latent_dist.sample(generator)
-    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
-        return encoder_output.latent_dist.mode()
-    elif hasattr(encoder_output, "latents"):
-        return encoder_output.latents
-    else:
-        raise AttributeError("Could not access latents of provided encoder_output")
-
-
-def calculate_dimensions(target_area, ratio):
-    width = math.sqrt(target_area * ratio)
-    height = width / ratio
-
-    width = round(width / 32) * 32
-    height = round(height / 32) * 32
-
-    return width, height
-
-
-class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
-    r"""
-    The Qwen-Image-Edit pipeline for image editing.
-
-    Args:
-        transformer ([`QwenImageTransformer2DModel`]):
-            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
-        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
-            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
-        vae ([`AutoencoderKL`]):
-            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
-        text_encoder ([`Qwen2.5-VL-7B-Instruct`]):
-            [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct), specifically the
-            [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) variant.
-        tokenizer (`QwenTokenizer`):
-            Tokenizer of class
-            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
-    """
-
-    model_cpu_offload_seq = "text_encoder->transformer->vae"
-    _callback_tensor_inputs = ["latents", "prompt_embeds"]
-
-    def __init__(
-        self,
-        scheduler: FlowMatchEulerDiscreteScheduler,
-        vae: AutoencoderKLQwenImage,
-        text_encoder: Qwen2_5_VLForConditionalGeneration,
-        tokenizer: Qwen2Tokenizer,
-        processor: Qwen2VLProcessor,
-        transformer: QwenImageTransformer2DModel,
-    ):
-        super().__init__()
-
-        self.register_modules(
-            vae=vae,
-            text_encoder=text_encoder,
-            tokenizer=tokenizer,
-            processor=processor,
-            transformer=transformer,
-            scheduler=scheduler,
-        )
-        self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
-        self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 16
-        # QwenImage latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible
-        # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
-        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
-        self.tokenizer_max_length = 1024
-
-        self.prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
-        self.prompt_template_encode_start_idx = 64
-        self.default_sample_size = 128
-
-    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._extract_masked_hidden
-    def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
-        bool_mask = mask.bool()
-        valid_lengths = bool_mask.sum(dim=1)
-        selected = hidden_states[bool_mask]
-        split_result = torch.split(selected, valid_lengths.tolist(), dim=0)
-
-        return split_result
-
-    def _get_qwen_prompt_embeds(
-        self,
-        prompt: Union[str, List[str]] = None,
-        image: Optional[torch.Tensor] = None,
-        device: Optional[torch.device] = None,
-        dtype: Optional[torch.dtype] = None,
-    ):
-        device = device or self._execution_device
-        dtype = dtype or self.text_encoder.dtype
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>"
-        if isinstance(image, list):
-            base_img_prompt = ""
-            for i, img in enumerate(image):
-                base_img_prompt += img_prompt_template.format(i + 1)
-        elif image is not None:
-            base_img_prompt = img_prompt_template.format(1)
-        else:
-            base_img_prompt = ""
-
-        template = self.prompt_template_encode
-
-        drop_idx = self.prompt_template_encode_start_idx
-        txt = [template.format(base_img_prompt + e) for e in prompt]
-
-        model_inputs = self.processor(
-            text=txt,
-            images=image,
-            padding=True,
-            return_tensors="pt",
-        ).to(device)
-
-        outputs = self.text_encoder(
-            input_ids=model_inputs.input_ids,
-            attention_mask=model_inputs.attention_mask,
-            pixel_values=model_inputs.pixel_values,
-            image_grid_thw=model_inputs.image_grid_thw,
-            output_hidden_states=True,
-        )
-
-        hidden_states = outputs.hidden_states[-1]
-        split_hidden_states = self._extract_masked_hidden(hidden_states, model_inputs.attention_mask)
-        split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
-        attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
-        max_seq_len = max([e.size(0) for e in split_hidden_states])
-        prompt_embeds = torch.stack(
-            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
-        )
-        encoder_attention_mask = torch.stack(
-            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
-        )
-
-        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
-
-        return prompt_embeds, encoder_attention_mask
-
-    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.encode_prompt
-    def encode_prompt(
-        self,
-        prompt: Union[str, List[str]],
-        image: Optional[torch.Tensor] = None,
-        device: Optional[torch.device] = None,
-        num_images_per_prompt: int = 1,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        prompt_embeds_mask: Optional[torch.Tensor] = None,
-        max_sequence_length: int = 1024,
-    ):
-        r"""
-
-        Args:
-            prompt (`str` or `List[str]`, *optional*):
-                prompt to be encoded
-            image (`torch.Tensor`, *optional*):
-                image to be encoded
-            device: (`torch.device`):
-                torch device
-            num_images_per_prompt (`int`):
-                number of images that should be generated per prompt
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-        """
-        device = device or self._execution_device
-
-        prompt = [prompt] if isinstance(prompt, str) else prompt
-        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
-
-        if prompt_embeds is None:
-            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image, device)
-
-        _, seq_len, _ = prompt_embeds.shape
-        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
-        prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
-        prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
-
-        return prompt_embeds, prompt_embeds_mask
-
-    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.check_inputs
-    def check_inputs(
-        self,
-        prompt,
-        height,
-        width,
-        negative_prompt=None,
-        prompt_embeds=None,
-        negative_prompt_embeds=None,
-        prompt_embeds_mask=None,
-        negative_prompt_embeds_mask=None,
-        callback_on_step_end_tensor_inputs=None,
-        max_sequence_length=None,
-    ):
-        if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0:
-            logger.warning(
-                f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly"
-            )
-
-        if callback_on_step_end_tensor_inputs is not None and not all(
-            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
-        ):
-            raise ValueError(
-                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
-            )
-
-        if prompt is not None and prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
-                " only forward one of the two."
-            )
-        elif prompt is None and prompt_embeds is None:
-            raise ValueError(
-                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
-            )
-        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
-            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
-
-        if negative_prompt is not None and negative_prompt_embeds is not None:
-            raise ValueError(
-                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
-                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
-            )
-
-        if prompt_embeds is not None and prompt_embeds_mask is None:
-            raise ValueError(
-                "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. Make sure to generate `prompt_embeds_mask` from the same text encoder that was used to generate `prompt_embeds`."
-            )
-        if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None:
-            raise ValueError(
-                "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. Make sure to generate `negative_prompt_embeds_mask` from the same text encoder that was used to generate `negative_prompt_embeds`."
-            )
-
-        if max_sequence_length is not None and max_sequence_length > 1024:
-            raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
-
-    @staticmethod
-    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._pack_latents
-    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
-        latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
-        latents = latents.permute(0, 2, 4, 1, 3, 5)
-        latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
-
-        return latents
-
-    @staticmethod
-    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._unpack_latents
-    def _unpack_latents(latents, height, width, vae_scale_factor):
-        batch_size, num_patches, channels = latents.shape
-
-        # VAE applies 8x compression on images but we must also account for packing which requires
-        # latent height and width to be divisible by 2.
-        height = 2 * (int(height) // (vae_scale_factor * 2))
-        width = 2 * (int(width) // (vae_scale_factor * 2))
-
-        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
-        latents = latents.permute(0, 3, 1, 4, 2, 5)
-
-        latents = latents.reshape(batch_size, channels // (2 * 2), 1, height, width)
-
-        return latents
-
-    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline._encode_vae_image
-    def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
-        if isinstance(generator, list):
-            image_latents = [
-                retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="argmax")
-                for i in range(image.shape[0])
-            ]
-            image_latents = torch.cat(image_latents, dim=0)
-        else:
-            image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax")
-        latents_mean = (
-            torch.tensor(self.vae.config.latents_mean)
-            .view(1, self.latent_channels, 1, 1, 1)
-            .to(image_latents.device, image_latents.dtype)
-        )
-        latents_std = (
-            torch.tensor(self.vae.config.latents_std)
-            .view(1, self.latent_channels, 1, 1, 1)
-            .to(image_latents.device, image_latents.dtype)
-        )
-        image_latents = (image_latents - latents_mean) / latents_std
-
-        return image_latents
-
-    def prepare_latents(
-        self,
-        images,
-        batch_size,
-        num_channels_latents,
-        height,
-        width,
-        dtype,
-        device,
-        generator,
-        latents=None,
-    ):
-        # VAE applies 8x compression on images but we must also account for packing which requires
-        # latent height and width to be divisible by 2.
-        height = 2 * (int(height) // (self.vae_scale_factor * 2))
-        width = 2 * (int(width) // (self.vae_scale_factor * 2))
-
-        shape = (batch_size, 1, num_channels_latents, height, width)
-
-        image_latents = None
-        if images is not None:
-            if not isinstance(images, list):
-                images = [images]
-            all_image_latents = []
-            for image in images:
-                image = image.to(device=device, dtype=dtype)
-                if image.shape[1] != self.latent_channels:
-                    image_latents = self._encode_vae_image(image=image, generator=generator)
-                else:
-                    image_latents = image
-                if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
-                    # expand init_latents for batch_size
-                    additional_image_per_prompt = batch_size // image_latents.shape[0]
-                    image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
-                elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
-                    raise ValueError(
-                        f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
-                    )
-                else:
-                    image_latents = torch.cat([image_latents], dim=0)
-
-                image_latent_height, image_latent_width = image_latents.shape[3:]
-                image_latents = self._pack_latents(
-                    image_latents, batch_size, num_channels_latents, image_latent_height, image_latent_width
-                )
-                all_image_latents.append(image_latents)
-            image_latents = torch.cat(all_image_latents, dim=1)
-
-        if isinstance(generator, list) and len(generator) != batch_size:
-            raise ValueError(
-                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
-            )
-        if latents is None:
-            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-            latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
-        else:
-            latents = latents.to(device=device, dtype=dtype)
-
-        return latents, image_latents
-
-    @property
-    def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def attention_kwargs(self):
-        return self._attention_kwargs
-
-    @property
-    def num_timesteps(self):
-        return self._num_timesteps
-
-    @property
-    def current_timestep(self):
-        return self._current_timestep
-
-    @property
-    def interrupt(self):
-        return self._interrupt
-
-    @torch.no_grad()
-    @replace_example_docstring(EXAMPLE_DOC_STRING)
-    def __call__(
-        self,
-        image: Optional[PipelineImageInput] = None,
-        prompt: Union[str, List[str]] = None,
-        negative_prompt: Union[str, List[str]] = None,
-        true_cfg_scale: float = 4.0,
-        height: Optional[int] = None,
-        width: Optional[int] = None,
-        num_inference_steps: int = 50,
-        sigmas: Optional[List[float]] = None,
-        guidance_scale: Optional[float] = None,
-        num_images_per_prompt: int = 1,
-        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
-        latents: Optional[torch.Tensor] = None,
-        prompt_embeds: Optional[torch.Tensor] = None,
-        prompt_embeds_mask: Optional[torch.Tensor] = None,
-        negative_prompt_embeds: Optional[torch.Tensor] = None,
-        negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
-        output_type: Optional[str] = "pil",
-        return_dict: bool = True,
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
-        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
-        max_sequence_length: int = 512,
-    ):
-        r"""
-        Function invoked when calling the pipeline for generation.
-
-        Args:
-            image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
-                `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
-                numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
-                or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
-                list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
-                latents as `image`, but if passing latents directly it is not encoded again.
-            prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
-                instead.
-            negative_prompt (`str` or `List[str]`, *optional*):
-                The prompt or prompts not to guide the image generation. If not defined, one has to pass
-                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
-                not greater than `1`).
-            true_cfg_scale (`float`, *optional*, defaults to 1.0):
-                true_cfg_scale (`float`, *optional*, defaults to 1.0): Guidance scale as defined in [Classifier-Free
-                Diffusion Guidance](https://huggingface.co/papers/2207.12598). `true_cfg_scale` is defined as `w` of
-                equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Classifier-free guidance is
-                enabled by setting `true_cfg_scale > 1` and a provided `negative_prompt`. Higher guidance scale
-                encourages to generate images that are closely linked to the text `prompt`, usually at the expense of
-                lower image quality.
-            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The height in pixels of the generated image. This is set to 1024 by default for the best results.
-            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
-                The width in pixels of the generated image. This is set to 1024 by default for the best results.
-            num_inference_steps (`int`, *optional*, defaults to 50):
-                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
-                expense of slower inference.
-            sigmas (`List[float]`, *optional*):
-                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
-                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
-                will be used.
-            guidance_scale (`float`, *optional*, defaults to None):
-                A guidance scale value for guidance distilled models. Unlike the traditional classifier-free guidance
-                where the guidance scale is applied during inference through noise prediction rescaling, guidance
-                distilled models take the guidance scale directly as an input parameter during forward pass. Guidance
-                scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images
-                that are closely linked to the text `prompt`, usually at the expense of lower image quality. This
-                parameter in the pipeline is there to support future guidance-distilled models when they come up. It is
-                ignored when not using guidance distilled models. To enable traditional classifier-free guidance,
-                please pass `true_cfg_scale > 1.0` and `negative_prompt` (even an empty negative prompt like " " should
-                enable classifier-free guidance computations).
-            num_images_per_prompt (`int`, *optional*, defaults to 1):
-                The number of images to generate per prompt.
-            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
-                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
-                to make generation deterministic.
-            latents (`torch.Tensor`, *optional*):
-                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
-                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
-                tensor will be generated by sampling using the supplied random `generator`.
-            prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
-                provided, text embeddings will be generated from `prompt` input argument.
-            negative_prompt_embeds (`torch.Tensor`, *optional*):
-                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
-                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
-                argument.
-            output_type (`str`, *optional*, defaults to `"pil"`):
-                The output format of the generate image. Choose between
-                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~pipelines.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple.
-            attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            callback_on_step_end (`Callable`, *optional*):
-                A function that calls at the end of each denoising steps during the inference. The function is called
-                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
-                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
-                `callback_on_step_end_tensor_inputs`.
-            callback_on_step_end_tensor_inputs (`List`, *optional*):
-                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
-                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
-                `._callback_tensor_inputs` attribute of your pipeline class.
-            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
-
-        Examples:
-
-        Returns:
-            [`~pipelines.qwenimage.QwenImagePipelineOutput`] or `tuple`:
-            [`~pipelines.qwenimage.QwenImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
-            returning a tuple, the first element is a list with the generated images.
-        """
-        image_size = image[0].size if isinstance(image, list) else image.size
-        calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1])
-        height = height or calculated_height
-        width = width or calculated_width
-
-        multiple_of = self.vae_scale_factor * 2
-        width = width // multiple_of * multiple_of
-        height = height // multiple_of * multiple_of
-
-        # 1. Check inputs. Raise error if not correct
-        self.check_inputs(
-            prompt,
-            height,
-            width,
-            negative_prompt=negative_prompt,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            prompt_embeds_mask=prompt_embeds_mask,
-            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
-            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
-            max_sequence_length=max_sequence_length,
-        )
-
-        self._guidance_scale = guidance_scale
-        self._attention_kwargs = attention_kwargs
-        self._current_timestep = None
-        self._interrupt = False
-
-        # 2. Define call parameters
-        if prompt is not None and isinstance(prompt, str):
-            batch_size = 1
-        elif prompt is not None and isinstance(prompt, list):
-            batch_size = len(prompt)
-        else:
-            batch_size = prompt_embeds.shape[0]
-
-        device = self._execution_device
-        # 3. Preprocess image
-        if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
-            if not isinstance(image, list):
-                image = [image]
-            condition_image_sizes = []
-            condition_images = []
-            vae_image_sizes = []
-            vae_images = []
-            for img in image:
-                image_width, image_height = img.size
-                condition_width, condition_height = calculate_dimensions(
-                    CONDITION_IMAGE_SIZE, image_width / image_height
-                )
-                vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height)
-                condition_image_sizes.append((condition_width, condition_height))
-                vae_image_sizes.append((vae_width, vae_height))
-                condition_images.append(self.image_processor.resize(img, condition_height, condition_width))
-                vae_images.append(self.image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2))
-
-        has_neg_prompt = negative_prompt is not None or (
-            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
-        )
-
-        if true_cfg_scale > 1 and not has_neg_prompt:
-            logger.warning(
-                f"true_cfg_scale is passed as {true_cfg_scale}, but classifier-free guidance is not enabled since no negative_prompt is provided."
-            )
-        elif true_cfg_scale <= 1 and has_neg_prompt:
-            logger.warning(
-                " negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1"
-            )
-
-        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
-        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
-            image=condition_images,
-            prompt=prompt,
-            prompt_embeds=prompt_embeds,
-            prompt_embeds_mask=prompt_embeds_mask,
-            device=device,
-            num_images_per_prompt=num_images_per_prompt,
-            max_sequence_length=max_sequence_length,
-        )
-        if do_true_cfg:
-            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
-                image=condition_images,
-                prompt=negative_prompt,
-                prompt_embeds=negative_prompt_embeds,
-                prompt_embeds_mask=negative_prompt_embeds_mask,
-                device=device,
-                num_images_per_prompt=num_images_per_prompt,
-                max_sequence_length=max_sequence_length,
-            )
-
-        # 4. Prepare latent variables
-        num_channels_latents = self.transformer.config.in_channels // 4
-        latents, image_latents = self.prepare_latents(
-            vae_images,
-            batch_size * num_images_per_prompt,
-            num_channels_latents,
-            height,
-            width,
-            prompt_embeds.dtype,
-            device,
-            generator,
-            latents,
-        )
-        img_shapes = [
-            [
-                (1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2),
-                *[
-                    (1, vae_height // self.vae_scale_factor // 2, vae_width // self.vae_scale_factor // 2)
-                    for vae_width, vae_height in vae_image_sizes
-                ],
-            ]
-        ] * batch_size
-
-        # 5. Prepare timesteps
-        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
-        image_seq_len = latents.shape[1]
-        mu = calculate_shift(
-            image_seq_len,
-            self.scheduler.config.get("base_image_seq_len", 256),
-            self.scheduler.config.get("max_image_seq_len", 4096),
-            self.scheduler.config.get("base_shift", 0.5),
-            self.scheduler.config.get("max_shift", 1.15),
-        )
-        timesteps, num_inference_steps = retrieve_timesteps(
-            self.scheduler,
-            num_inference_steps,
-            device,
-            sigmas=sigmas,
-            mu=mu,
-        )
-        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
-        self._num_timesteps = len(timesteps)
-
-        # handle guidance
-        if self.transformer.config.guidance_embeds and guidance_scale is None:
-            raise ValueError("guidance_scale is required for guidance-distilled model.")
-        elif self.transformer.config.guidance_embeds:
-            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
-            guidance = guidance.expand(latents.shape[0])
-        elif not self.transformer.config.guidance_embeds and guidance_scale is not None:
-            logger.warning(
-                f"guidance_scale is passed as {guidance_scale}, but ignored since the model is not guidance-distilled."
-            )
-            guidance = None
-        elif not self.transformer.config.guidance_embeds and guidance_scale is None:
-            guidance = None
-
-        if self.attention_kwargs is None:
-            self._attention_kwargs = {}
-
-        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
-        
-        image_rotary_emb = self.transformer.pos_embed(img_shapes, txt_seq_lens, device=latents.device)
-        if do_true_cfg:
-            negative_txt_seq_lens = (
-                negative_prompt_embeds_mask.sum(dim=1).tolist()
-                if negative_prompt_embeds_mask is not None
-                else None
-            )
-            uncond_image_rotary_emb = self.transformer.pos_embed(
-                img_shapes, negative_txt_seq_lens, device=latents.device
-            )
-        else:
-            uncond_image_rotary_emb = None
-
-        # 6. Denoising loop
-        self.scheduler.set_begin_index(0)
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                if self.interrupt:
-                    continue
-
-                self._current_timestep = t
-
-                latent_model_input = latents
-                if image_latents is not None:
-                    latent_model_input = torch.cat([latents, image_latents], dim=1)
-
-                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                timestep = t.expand(latents.shape[0]).to(latents.dtype)
-                with self.transformer.cache_context("cond"):
-                    noise_pred = self.transformer(
-                        hidden_states=latent_model_input,
-                        timestep=timestep / 1000,
-                        guidance=guidance,
-                        encoder_hidden_states_mask=prompt_embeds_mask,
-                        encoder_hidden_states=prompt_embeds,
-                        image_rotary_emb=image_rotary_emb,
-                        attention_kwargs=self.attention_kwargs,
-                        return_dict=False,
-                    )[0]
-                    noise_pred = noise_pred[:, : latents.size(1)]
-
-                if do_true_cfg:
-                    with self.transformer.cache_context("uncond"):
-                        neg_noise_pred = self.transformer(
-                            hidden_states=latent_model_input,
-                            timestep=timestep / 1000,
-                            guidance=guidance,
-                            encoder_hidden_states_mask=negative_prompt_embeds_mask,
-                            encoder_hidden_states=negative_prompt_embeds,
-                            image_rotary_emb=uncond_image_rotary_emb,
-                            attention_kwargs=self.attention_kwargs,
-                            return_dict=False,
-                        )[0]
-                    neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
-                    comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
-
-                    cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
-                    noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
-                    noise_pred = comb_pred * (cond_norm / noise_norm)
-
-                # compute the previous noisy sample x_t -> x_t-1
-                latents_dtype = latents.dtype
-                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-
-                if latents.dtype != latents_dtype:
-                    if torch.backends.mps.is_available():
-                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
-                        latents = latents.to(latents_dtype)
-
-                if callback_on_step_end is not None:
-                    callback_kwargs = {}
-                    for k in callback_on_step_end_tensor_inputs:
-                        callback_kwargs[k] = locals()[k]
-                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
-
-                    latents = callback_outputs.pop("latents", latents)
-                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
-
-                # call the callback, if provided
-                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
-                    progress_bar.update()
-
-                if XLA_AVAILABLE:
-                    xm.mark_step()
-
-        self._current_timestep = None
-        if output_type == "latent":
-            image = latents
-        else:
-            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
-            latents = latents.to(self.vae.dtype)
-            latents_mean = (
-                torch.tensor(self.vae.config.latents_mean)
-                .view(1, self.vae.config.z_dim, 1, 1, 1)
-                .to(latents.device, latents.dtype)
-            )
-            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
-                latents.device, latents.dtype
-            )
-            latents = latents / latents_std + latents_mean
-            image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
-            image = self.image_processor.postprocess(image, output_type=output_type)
-
-        # Offload all models
-        self.maybe_free_model_hooks()
-
-        if not return_dict:
-            return (image,)
-
-        return QwenImagePipelineOutput(images=image)
+# Copyright 2025 Qwen-Image Team and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+import math
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor
+
+from diffusers.image_processor import PipelineImageInput, VaeImageProcessor
+from diffusers.loaders import QwenImageLoraLoaderMixin
+from diffusers.models import AutoencoderKLQwenImage, QwenImageTransformer2DModel
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import is_torch_xla_available, logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.pipelines.qwenimage.pipeline_output import QwenImagePipelineOutput
+
+
+if is_torch_xla_available():
+    import torch_xla.core.xla_model as xm
+
+    XLA_AVAILABLE = True
+else:
+    XLA_AVAILABLE = False
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from PIL import Image
+        >>> from diffusers import QwenImageEditPlusPipeline
+        >>> from diffusers.utils import load_image
+
+        >>> pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2511", torch_dtype=torch.bfloat16)
+        >>> pipe.to("cuda")
+        >>> image = load_image("input.png").convert("RGB")
+        >>> prompt = (
+        ...     "Make Pikachu hold a sign that says 'Qwen Edit is awesome', yarn art style, detailed, vibrant colors"
+        ... )
+        >>> # Depending on the variant being used, the pipeline call will slightly vary.
+        >>> # Refer to the pipeline documentation for more details.
+        >>> image = pipe(image, prompt, num_inference_steps=50).images[0]
+        >>> image.save("qwenimage_edit_plus.png")
+        ```
+"""
+
+CONDITION_IMAGE_SIZE = 384 * 384
+VAE_IMAGE_SIZE = 1024 * 1024
+
+
+# Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.calculate_shift
+def calculate_shift(
+    image_seq_len,
+    base_seq_len: int = 256,
+    max_seq_len: int = 4096,
+    base_shift: float = 0.5,
+    max_shift: float = 1.15,
+):
+    m = (max_shift - base_shift) / (max_seq_len - base_seq_len)
+    b = base_shift - m * base_seq_len
+    mu = image_seq_len * m + b
+    return mu
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents
+def retrieve_latents(
+    encoder_output: torch.Tensor, generator: Optional[torch.Generator] = None, sample_mode: str = "sample"
+):
+    if hasattr(encoder_output, "latent_dist") and sample_mode == "sample":
+        return encoder_output.latent_dist.sample(generator)
+    elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax":
+        return encoder_output.latent_dist.mode()
+    elif hasattr(encoder_output, "latents"):
+        return encoder_output.latents
+    else:
+        raise AttributeError("Could not access latents of provided encoder_output")
+
+
+def calculate_dimensions(target_area, ratio):
+    width = math.sqrt(target_area * ratio)
+    height = width / ratio
+
+    width = round(width / 32) * 32
+    height = round(height / 32) * 32
+
+    return width, height
+
+
+class QwenImageEditPlusPipeline(DiffusionPipeline, QwenImageLoraLoaderMixin):
+    r"""
+    The Qwen-Image-Edit pipeline for image editing.
+
+    Args:
+        transformer ([`QwenImageTransformer2DModel`]):
+            Conditional Transformer (MMDiT) architecture to denoise the encoded image latents.
+        scheduler ([`FlowMatchEulerDiscreteScheduler`]):
+            A scheduler to be used in combination with `transformer` to denoise the encoded image latents.
+        vae ([`AutoencoderKL`]):
+            Variational Auto-Encoder (VAE) Model to encode and decode images to and from latent representations.
+        text_encoder ([`Qwen2.5-VL-7B-Instruct`]):
+            [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct), specifically the
+            [Qwen2.5-VL-7B-Instruct](https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct) variant.
+        tokenizer (`QwenTokenizer`):
+            Tokenizer of class
+            [CLIPTokenizer](https://huggingface.co/docs/transformers/en/model_doc/clip#transformers.CLIPTokenizer).
+    """
+
+    model_cpu_offload_seq = "text_encoder->transformer->vae"
+    _callback_tensor_inputs = ["latents", "prompt_embeds"]
+
+    def __init__(
+        self,
+        scheduler: FlowMatchEulerDiscreteScheduler,
+        vae: AutoencoderKLQwenImage,
+        text_encoder: Qwen2_5_VLForConditionalGeneration,
+        tokenizer: Qwen2Tokenizer,
+        processor: Qwen2VLProcessor,
+        transformer: QwenImageTransformer2DModel,
+    ):
+        super().__init__()
+
+        self.register_modules(
+            vae=vae,
+            text_encoder=text_encoder,
+            tokenizer=tokenizer,
+            processor=processor,
+            transformer=transformer,
+            scheduler=scheduler,
+        )
+        self.vae_scale_factor = 2 ** len(self.vae.temperal_downsample) if getattr(self, "vae", None) else 8
+        self.latent_channels = self.vae.config.z_dim if getattr(self, "vae", None) else 16
+        # QwenImage latents are turned into 2x2 patches and packed. This means the latent width and height has to be divisible
+        # by the patch size. So the vae scale factor is multiplied by the patch size to account for this
+        self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor * 2)
+        self.tokenizer_max_length = 1024
+
+        self.prompt_template_encode = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+        self.prompt_template_encode_start_idx = 64
+        self.default_sample_size = 128
+
+    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._extract_masked_hidden
+    def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor):
+        bool_mask = mask.bool()
+        valid_lengths = bool_mask.sum(dim=1)
+        selected = hidden_states[bool_mask]
+        split_result = torch.split(selected, valid_lengths.tolist(), dim=0)
+
+        return split_result
+
+    def _get_qwen_prompt_embeds(
+        self,
+        prompt: Union[str, List[str]] = None,
+        image: Optional[torch.Tensor] = None,
+        device: Optional[torch.device] = None,
+        dtype: Optional[torch.dtype] = None,
+    ):
+        device = device or self._execution_device
+        dtype = dtype or self.text_encoder.dtype
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        img_prompt_template = "Picture {}: <|vision_start|><|image_pad|><|vision_end|>"
+        if isinstance(image, list):
+            base_img_prompt = ""
+            for i, img in enumerate(image):
+                base_img_prompt += img_prompt_template.format(i + 1)
+        elif image is not None:
+            base_img_prompt = img_prompt_template.format(1)
+        else:
+            base_img_prompt = ""
+
+        template = self.prompt_template_encode
+
+        drop_idx = self.prompt_template_encode_start_idx
+        txt = [template.format(base_img_prompt + e) for e in prompt]
+
+        model_inputs = self.processor(
+            text=txt,
+            images=image,
+            padding=True,
+            return_tensors="pt",
+        ).to(device)
+
+        outputs = self.text_encoder(
+            input_ids=model_inputs.input_ids,
+            attention_mask=model_inputs.attention_mask,
+            pixel_values=model_inputs.pixel_values,
+            image_grid_thw=model_inputs.image_grid_thw,
+            output_hidden_states=True,
+        )
+
+        hidden_states = outputs.hidden_states[-1]
+        split_hidden_states = self._extract_masked_hidden(hidden_states, model_inputs.attention_mask)
+        split_hidden_states = [e[drop_idx:] for e in split_hidden_states]
+        attn_mask_list = [torch.ones(e.size(0), dtype=torch.long, device=e.device) for e in split_hidden_states]
+        max_seq_len = max([e.size(0) for e in split_hidden_states])
+        prompt_embeds = torch.stack(
+            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0), u.size(1))]) for u in split_hidden_states]
+        )
+        encoder_attention_mask = torch.stack(
+            [torch.cat([u, u.new_zeros(max_seq_len - u.size(0))]) for u in attn_mask_list]
+        )
+
+        prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+        return prompt_embeds, encoder_attention_mask
+
+    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.encode_prompt
+    def encode_prompt(
+        self,
+        prompt: Union[str, List[str]],
+        image: Optional[torch.Tensor] = None,
+        device: Optional[torch.device] = None,
+        num_images_per_prompt: int = 1,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_embeds_mask: Optional[torch.Tensor] = None,
+        max_sequence_length: int = 1024,
+    ):
+        r"""
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                prompt to be encoded
+            image (`torch.Tensor`, *optional*):
+                image to be encoded
+            device: (`torch.device`):
+                torch device
+            num_images_per_prompt (`int`):
+                number of images that should be generated per prompt
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+        """
+        device = device or self._execution_device
+
+        prompt = [prompt] if isinstance(prompt, str) else prompt
+        batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0]
+
+        if prompt_embeds is None:
+            prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image, device)
+
+        _, seq_len, _ = prompt_embeds.shape
+        prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1)
+        prompt_embeds_mask = prompt_embeds_mask.repeat(1, num_images_per_prompt, 1)
+        prompt_embeds_mask = prompt_embeds_mask.view(batch_size * num_images_per_prompt, seq_len)
+
+        return prompt_embeds, prompt_embeds_mask
+
+    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline.check_inputs
+    def check_inputs(
+        self,
+        prompt,
+        height,
+        width,
+        negative_prompt=None,
+        prompt_embeds=None,
+        negative_prompt_embeds=None,
+        prompt_embeds_mask=None,
+        negative_prompt_embeds_mask=None,
+        callback_on_step_end_tensor_inputs=None,
+        max_sequence_length=None,
+    ):
+        if height % (self.vae_scale_factor * 2) != 0 or width % (self.vae_scale_factor * 2) != 0:
+            logger.warning(
+                f"`height` and `width` have to be divisible by {self.vae_scale_factor * 2} but are {height} and {width}. Dimensions will be resized accordingly"
+            )
+
+        if callback_on_step_end_tensor_inputs is not None and not all(
+            k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+        ):
+            raise ValueError(
+                f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+            )
+
+        if prompt is not None and prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+                " only forward one of the two."
+            )
+        elif prompt is None and prompt_embeds is None:
+            raise ValueError(
+                "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+            )
+        elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+            raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+        if negative_prompt is not None and negative_prompt_embeds is not None:
+            raise ValueError(
+                f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+                f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+            )
+
+        if prompt_embeds is not None and prompt_embeds_mask is None:
+            raise ValueError(
+                "If `prompt_embeds` are provided, `prompt_embeds_mask` also have to be passed. Make sure to generate `prompt_embeds_mask` from the same text encoder that was used to generate `prompt_embeds`."
+            )
+        if negative_prompt_embeds is not None and negative_prompt_embeds_mask is None:
+            raise ValueError(
+                "If `negative_prompt_embeds` are provided, `negative_prompt_embeds_mask` also have to be passed. Make sure to generate `negative_prompt_embeds_mask` from the same text encoder that was used to generate `negative_prompt_embeds`."
+            )
+
+        if max_sequence_length is not None and max_sequence_length > 1024:
+            raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}")
+
+    @staticmethod
+    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._pack_latents
+    def _pack_latents(latents, batch_size, num_channels_latents, height, width):
+        latents = latents.view(batch_size, num_channels_latents, height // 2, 2, width // 2, 2)
+        latents = latents.permute(0, 2, 4, 1, 3, 5)
+        latents = latents.reshape(batch_size, (height // 2) * (width // 2), num_channels_latents * 4)
+
+        return latents
+
+    @staticmethod
+    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage.QwenImagePipeline._unpack_latents
+    def _unpack_latents(latents, height, width, vae_scale_factor):
+        batch_size, num_patches, channels = latents.shape
+
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (vae_scale_factor * 2))
+        width = 2 * (int(width) // (vae_scale_factor * 2))
+
+        latents = latents.view(batch_size, height // 2, width // 2, channels // 4, 2, 2)
+        latents = latents.permute(0, 3, 1, 4, 2, 5)
+
+        latents = latents.reshape(batch_size, channels // (2 * 2), 1, height, width)
+
+        return latents
+
+    # Copied from diffusers.pipelines.qwenimage.pipeline_qwenimage_edit.QwenImageEditPipeline._encode_vae_image
+    def _encode_vae_image(self, image: torch.Tensor, generator: torch.Generator):
+        if isinstance(generator, list):
+            image_latents = [
+                retrieve_latents(self.vae.encode(image[i : i + 1]), generator=generator[i], sample_mode="argmax")
+                for i in range(image.shape[0])
+            ]
+            image_latents = torch.cat(image_latents, dim=0)
+        else:
+            image_latents = retrieve_latents(self.vae.encode(image), generator=generator, sample_mode="argmax")
+        latents_mean = (
+            torch.tensor(self.vae.config.latents_mean)
+            .view(1, self.latent_channels, 1, 1, 1)
+            .to(image_latents.device, image_latents.dtype)
+        )
+        latents_std = (
+            torch.tensor(self.vae.config.latents_std)
+            .view(1, self.latent_channels, 1, 1, 1)
+            .to(image_latents.device, image_latents.dtype)
+        )
+        image_latents = (image_latents - latents_mean) / latents_std
+
+        return image_latents
+
+    def prepare_latents(
+        self,
+        images,
+        batch_size,
+        num_channels_latents,
+        height,
+        width,
+        dtype,
+        device,
+        generator,
+        latents=None,
+    ):
+        # VAE applies 8x compression on images but we must also account for packing which requires
+        # latent height and width to be divisible by 2.
+        height = 2 * (int(height) // (self.vae_scale_factor * 2))
+        width = 2 * (int(width) // (self.vae_scale_factor * 2))
+
+        shape = (batch_size, 1, num_channels_latents, height, width)
+
+        image_latents = None
+        if images is not None:
+            if not isinstance(images, list):
+                images = [images]
+            all_image_latents = []
+            for image in images:
+                image = image.to(device=device, dtype=dtype)
+                if image.shape[1] != self.latent_channels:
+                    image_latents = self._encode_vae_image(image=image, generator=generator)
+                else:
+                    image_latents = image
+                if batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] == 0:
+                    # expand init_latents for batch_size
+                    additional_image_per_prompt = batch_size // image_latents.shape[0]
+                    image_latents = torch.cat([image_latents] * additional_image_per_prompt, dim=0)
+                elif batch_size > image_latents.shape[0] and batch_size % image_latents.shape[0] != 0:
+                    raise ValueError(
+                        f"Cannot duplicate `image` of batch size {image_latents.shape[0]} to {batch_size} text prompts."
+                    )
+                else:
+                    image_latents = torch.cat([image_latents], dim=0)
+
+                image_latent_height, image_latent_width = image_latents.shape[3:]
+                image_latents = self._pack_latents(
+                    image_latents, batch_size, num_channels_latents, image_latent_height, image_latent_width
+                )
+                all_image_latents.append(image_latents)
+            image_latents = torch.cat(all_image_latents, dim=1)
+
+        if isinstance(generator, list) and len(generator) != batch_size:
+            raise ValueError(
+                f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+                f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+            )
+        if latents is None:
+            latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+            latents = self._pack_latents(latents, batch_size, num_channels_latents, height, width)
+        else:
+            latents = latents.to(device=device, dtype=dtype)
+
+        return latents, image_latents
+
+    @property
+    def guidance_scale(self):
+        return self._guidance_scale
+
+    @property
+    def attention_kwargs(self):
+        return self._attention_kwargs
+
+    @property
+    def num_timesteps(self):
+        return self._num_timesteps
+
+    @property
+    def current_timestep(self):
+        return self._current_timestep
+
+    @property
+    def interrupt(self):
+        return self._interrupt
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        image: Optional[PipelineImageInput] = None,
+        prompt: Union[str, List[str]] = None,
+        negative_prompt: Union[str, List[str]] = None,
+        true_cfg_scale: float = 4.0,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 50,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: Optional[float] = None,
+        num_images_per_prompt: int = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.Tensor] = None,
+        prompt_embeds: Optional[torch.Tensor] = None,
+        prompt_embeds_mask: Optional[torch.Tensor] = None,
+        negative_prompt_embeds: Optional[torch.Tensor] = None,
+        negative_prompt_embeds_mask: Optional[torch.Tensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            image (`torch.Tensor`, `PIL.Image.Image`, `np.ndarray`, `List[torch.Tensor]`, `List[PIL.Image.Image]`, or `List[np.ndarray]`):
+                `Image`, numpy array or tensor representing an image batch to be used as the starting point. For both
+                numpy array and pytorch tensor, the expected value range is between `[0, 1]` If it's a tensor or a list
+                or tensors, the expected shape should be `(B, C, H, W)` or `(C, H, W)`. If it is a numpy array or a
+                list of arrays, the expected shape should be `(B, H, W, C)` or `(H, W, C)` It can also accept image
+                latents as `image`, but if passing latents directly it is not encoded again.
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            negative_prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts not to guide the image generation. If not defined, one has to pass
+                `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `true_cfg_scale` is
+                not greater than `1`).
+            true_cfg_scale (`float`, *optional*, defaults to 1.0):
+                true_cfg_scale (`float`, *optional*, defaults to 1.0): Guidance scale as defined in [Classifier-Free
+                Diffusion Guidance](https://huggingface.co/papers/2207.12598). `true_cfg_scale` is defined as `w` of
+                equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Classifier-free guidance is
+                enabled by setting `true_cfg_scale > 1` and a provided `negative_prompt`. Higher guidance scale
+                encourages to generate images that are closely linked to the text `prompt`, usually at the expense of
+                lower image quality.
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, *optional*, defaults to None):
+                A guidance scale value for guidance distilled models. Unlike the traditional classifier-free guidance
+                where the guidance scale is applied during inference through noise prediction rescaling, guidance
+                distilled models take the guidance scale directly as an input parameter during forward pass. Guidance
+                scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images
+                that are closely linked to the text `prompt`, usually at the expense of lower image quality. This
+                parameter in the pipeline is there to support future guidance-distilled models when they come up. It is
+                ignored when not using guidance distilled models. To enable traditional classifier-free guidance,
+                please pass `true_cfg_scale > 1.0` and `negative_prompt` (even an empty negative prompt like " " should
+                enable classifier-free guidance computations).
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.Tensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will be generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            negative_prompt_embeds (`torch.Tensor`, *optional*):
+                Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+                weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+                argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.qwenimage.QwenImagePipelineOutput`] instead of a plain tuple.
+            attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.qwenimage.QwenImagePipelineOutput`] or `tuple`:
+            [`~pipelines.qwenimage.QwenImagePipelineOutput`] if `return_dict` is True, otherwise a `tuple`. When
+            returning a tuple, the first element is a list with the generated images.
+        """
+        image_size = image[0].size if isinstance(image, list) else image.size
+        calculated_width, calculated_height = calculate_dimensions(1024 * 1024, image_size[0] / image_size[1])
+        height = height or calculated_height
+        width = width or calculated_width
+
+        multiple_of = self.vae_scale_factor * 2
+        width = width // multiple_of * multiple_of
+        height = height // multiple_of * multiple_of
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            height,
+            width,
+            negative_prompt=negative_prompt,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            negative_prompt_embeds_mask=negative_prompt_embeds_mask,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._attention_kwargs = attention_kwargs
+        self._current_timestep = None
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        # 3. Preprocess image
+        if image is not None and not (isinstance(image, torch.Tensor) and image.size(1) == self.latent_channels):
+            if not isinstance(image, list):
+                image = [image]
+            condition_image_sizes = []
+            condition_images = []
+            vae_image_sizes = []
+            vae_images = []
+            for img in image:
+                image_width, image_height = img.size
+                condition_width, condition_height = calculate_dimensions(
+                    CONDITION_IMAGE_SIZE, image_width / image_height
+                )
+                vae_width, vae_height = calculate_dimensions(VAE_IMAGE_SIZE, image_width / image_height)
+                condition_image_sizes.append((condition_width, condition_height))
+                vae_image_sizes.append((vae_width, vae_height))
+                condition_images.append(self.image_processor.resize(img, condition_height, condition_width))
+                vae_images.append(self.image_processor.preprocess(img, vae_height, vae_width).unsqueeze(2))
+
+        has_neg_prompt = negative_prompt is not None or (
+            negative_prompt_embeds is not None and negative_prompt_embeds_mask is not None
+        )
+
+        if true_cfg_scale > 1 and not has_neg_prompt:
+            logger.warning(
+                f"true_cfg_scale is passed as {true_cfg_scale}, but classifier-free guidance is not enabled since no negative_prompt is provided."
+            )
+        elif true_cfg_scale <= 1 and has_neg_prompt:
+            logger.warning(
+                " negative_prompt is passed but classifier-free guidance is not enabled since true_cfg_scale <= 1"
+            )
+
+        do_true_cfg = true_cfg_scale > 1 and has_neg_prompt
+        prompt_embeds, prompt_embeds_mask = self.encode_prompt(
+            image=condition_images,
+            prompt=prompt,
+            prompt_embeds=prompt_embeds,
+            prompt_embeds_mask=prompt_embeds_mask,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+        )
+        if do_true_cfg:
+            negative_prompt_embeds, negative_prompt_embeds_mask = self.encode_prompt(
+                image=condition_images,
+                prompt=negative_prompt,
+                prompt_embeds=negative_prompt_embeds,
+                prompt_embeds_mask=negative_prompt_embeds_mask,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+            )
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, image_latents = self.prepare_latents(
+            vae_images,
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+        img_shapes = [
+            [
+                (1, height // self.vae_scale_factor // 2, width // self.vae_scale_factor // 2),
+                *[
+                    (1, vae_height // self.vae_scale_factor // 2, vae_width // self.vae_scale_factor // 2)
+                    for vae_width, vae_height in vae_image_sizes
+                ],
+            ]
+        ] * batch_size
+
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.get("base_image_seq_len", 256),
+            self.scheduler.config.get("max_image_seq_len", 4096),
+            self.scheduler.config.get("base_shift", 0.5),
+            self.scheduler.config.get("max_shift", 1.15),
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        if self.transformer.config.guidance_embeds and guidance_scale is None:
+            raise ValueError("guidance_scale is required for guidance-distilled model.")
+        elif self.transformer.config.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        elif not self.transformer.config.guidance_embeds and guidance_scale is not None:
+            logger.warning(
+                f"guidance_scale is passed as {guidance_scale}, but ignored since the model is not guidance-distilled."
+            )
+            guidance = None
+        elif not self.transformer.config.guidance_embeds and guidance_scale is None:
+            guidance = None
+
+        if self.attention_kwargs is None:
+            self._attention_kwargs = {}
+
+        txt_seq_lens = prompt_embeds_mask.sum(dim=1).tolist() if prompt_embeds_mask is not None else None
+        
+        image_rotary_emb = self.transformer.pos_embed(img_shapes, txt_seq_lens, device=latents.device)
+        if do_true_cfg:
+            negative_txt_seq_lens = (
+                negative_prompt_embeds_mask.sum(dim=1).tolist()
+                if negative_prompt_embeds_mask is not None
+                else None
+            )
+            uncond_image_rotary_emb = self.transformer.pos_embed(
+                img_shapes, negative_txt_seq_lens, device=latents.device
+            )
+        else:
+            uncond_image_rotary_emb = None
+
+        # 6. Denoising loop
+        self.scheduler.set_begin_index(0)
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                self._current_timestep = t
+
+                latent_model_input = latents
+                if image_latents is not None:
+                    latent_model_input = torch.cat([latents, image_latents], dim=1)
+
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latents.shape[0]).to(latents.dtype)
+                with self.transformer.cache_context("cond"):
+                    noise_pred = self.transformer(
+                        hidden_states=latent_model_input,
+                        timestep=timestep / 1000,
+                        guidance=guidance,
+                        encoder_hidden_states_mask=prompt_embeds_mask,
+                        encoder_hidden_states=prompt_embeds,
+                        image_rotary_emb=image_rotary_emb,
+                        attention_kwargs=self.attention_kwargs,
+                        return_dict=False,
+                    )[0]
+                    noise_pred = noise_pred[:, : latents.size(1)]
+
+                if do_true_cfg:
+                    with self.transformer.cache_context("uncond"):
+                        neg_noise_pred = self.transformer(
+                            hidden_states=latent_model_input,
+                            timestep=timestep / 1000,
+                            guidance=guidance,
+                            encoder_hidden_states_mask=negative_prompt_embeds_mask,
+                            encoder_hidden_states=negative_prompt_embeds,
+                            image_rotary_emb=uncond_image_rotary_emb,
+                            attention_kwargs=self.attention_kwargs,
+                            return_dict=False,
+                        )[0]
+                    neg_noise_pred = neg_noise_pred[:, : latents.size(1)]
+                    comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+
+                    cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
+                    noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
+                    noise_pred = comb_pred * (cond_norm / noise_norm)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        self._current_timestep = None
+        if output_type == "latent":
+            image = latents
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = latents.to(self.vae.dtype)
+            latents_mean = (
+                torch.tensor(self.vae.config.latents_mean)
+                .view(1, self.vae.config.z_dim, 1, 1, 1)
+                .to(latents.device, latents.dtype)
+            )
+            latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
+                latents.device, latents.dtype
+            )
+            latents = latents / latents_std + latents_mean
+            image = self.vae.decode(latents, return_dict=False)[0][:, :, 0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return QwenImagePipelineOutput(images=image)

qwenimage/qwen_fa3_processor.py

@@ -1,110 +1,110 @@
-"""
-Paired with a good language model. Thanks!
-"""
-
-import torch
-import torch.nn.functional as F
-from torch.nn.attention import sdpa_kernel, SDPBackend
-from typing import Optional, Tuple
-from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
-
-
-class QwenDoubleStreamAttnProcessorFA3:
-    """
-    FA3-grade attention processor for Qwen double-stream architecture.
-    Routes through PyTorch's cuDNN attention backend, which on Hopper (SM 9.0+)
-    dispatches to the same FlashAttention-3 family of kernels as vLLM's FA3 —
-    bundled with the base image so there's no external-kernel ABI risk.
-
-    Notes / limitations:
-    - General attention masks are not supported here. `is_causal=False` and no arbitrary mask.
-    - Expects an available `apply_rotary_emb_qwen` in scope (same as your non-FA3 processor).
-    """
-
-    _attention_backend = "fa3"  # for parity with your other processors, not used internally
-
-    @torch.no_grad()
-    def __call__(
-        self,
-        attn,  # Attention module with to_q/to_k/to_v/add_*_proj, norms, to_out, to_add_out, and .heads
-        hidden_states: torch.FloatTensor,                 # (B, S_img, D_model)  image stream
-        encoder_hidden_states: torch.FloatTensor = None,  # (B, S_txt, D_model)  text stream
-        encoder_hidden_states_mask: torch.FloatTensor = None,  # unused in FA3 path
-        attention_mask: Optional[torch.FloatTensor] = None,    # unused in FA3 path
-        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # (img_freqs, txt_freqs)
-    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
-        if encoder_hidden_states is None:
-            raise ValueError("QwenDoubleStreamAttnProcessorFA3 requires encoder_hidden_states (text stream).")
-        if attention_mask is not None:
-            raise NotImplementedError("attention_mask is not supported in this FA3 implementation.")
-
-        B, S_img, _ = hidden_states.shape
-        S_txt = encoder_hidden_states.shape[1]
-
-        # ---- QKV projections (image/sample stream) ----
-        img_q = attn.to_q(hidden_states)   # (B, S_img, D)
-        img_k = attn.to_k(hidden_states)
-        img_v = attn.to_v(hidden_states)
-
-        # ---- QKV projections (text/context stream) ----
-        txt_q = attn.add_q_proj(encoder_hidden_states)  # (B, S_txt, D)
-        txt_k = attn.add_k_proj(encoder_hidden_states)
-        txt_v = attn.add_v_proj(encoder_hidden_states)
-
-        # ---- Reshape to (B, S, H, D_h) ----
-        H = attn.heads
-        img_q = img_q.unflatten(-1, (H, -1))
-        img_k = img_k.unflatten(-1, (H, -1))
-        img_v = img_v.unflatten(-1, (H, -1))
-
-        txt_q = txt_q.unflatten(-1, (H, -1))
-        txt_k = txt_k.unflatten(-1, (H, -1))
-        txt_v = txt_v.unflatten(-1, (H, -1))
-
-        # ---- Q/K normalization (per your module contract) ----
-        if getattr(attn, "norm_q", None) is not None:
-            img_q = attn.norm_q(img_q)
-        if getattr(attn, "norm_k", None) is not None:
-            img_k = attn.norm_k(img_k)
-        if getattr(attn, "norm_added_q", None) is not None:
-            txt_q = attn.norm_added_q(txt_q)
-        if getattr(attn, "norm_added_k", None) is not None:
-            txt_k = attn.norm_added_k(txt_k)
-
-        # ---- RoPE (Qwen variant) ----
-        if image_rotary_emb is not None:
-            img_freqs, txt_freqs = image_rotary_emb
-            img_q = apply_rotary_emb_qwen(img_q, img_freqs, use_real=False)
-            img_k = apply_rotary_emb_qwen(img_k, img_freqs, use_real=False)
-            txt_q = apply_rotary_emb_qwen(txt_q, txt_freqs, use_real=False)
-            txt_k = apply_rotary_emb_qwen(txt_k, txt_freqs, use_real=False)
-
-        # ---- Joint attention over [text, image] along sequence axis ----
-        q = torch.cat([txt_q, img_q], dim=1)  # (B, S_total, H, D_h)
-        k = torch.cat([txt_k, img_k], dim=1)
-        v = torch.cat([txt_v, img_v], dim=1)
-
-        # SDPA wants (B, H, S, D_h); route through cuDNN's FA3 path on Hopper.
-        with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
-            out = F.scaled_dot_product_attention(
-                q.transpose(1, 2),
-                k.transpose(1, 2),
-                v.transpose(1, 2),
-                is_causal=False,
-            ).transpose(1, 2)  # back to (B, S_total, H, D_h)
-
-        # ---- Back to (B, S, D_model) ----
-        out = out.flatten(2, 3).to(q.dtype)
-
-        # Split back to text / image segments
-        txt_attn_out = out[:, :S_txt, :]
-        img_attn_out = out[:, S_txt:, :]
-
-        # ---- Output projections ----
-        img_attn_out = attn.to_out[0](img_attn_out)
-        if len(attn.to_out) > 1:
-            img_attn_out = attn.to_out[1](img_attn_out)  # dropout if present
-
-        txt_attn_out = attn.to_add_out(txt_attn_out)
-
-        return img_attn_out, txt_attn_out
+"""
+Paired with a good language model. Thanks!
+"""
+
+import torch
+import torch.nn.functional as F
+from torch.nn.attention import sdpa_kernel, SDPBackend
+from typing import Optional, Tuple
+from diffusers.models.transformers.transformer_qwenimage import apply_rotary_emb_qwen
+
+
+class QwenDoubleStreamAttnProcessorFA3:
+    """
+    FA3-grade attention processor for Qwen double-stream architecture.
+    Routes through PyTorch's cuDNN attention backend, which on Hopper (SM 9.0+)
+    dispatches to the same FlashAttention-3 family of kernels as vLLM's FA3 —
+    bundled with the base image so there's no external-kernel ABI risk.
+
+    Notes / limitations:
+    - General attention masks are not supported here. `is_causal=False` and no arbitrary mask.
+    - Expects an available `apply_rotary_emb_qwen` in scope (same as your non-FA3 processor).
+    """
+
+    _attention_backend = "fa3"  # for parity with your other processors, not used internally
+
+    @torch.no_grad()
+    def __call__(
+        self,
+        attn,  # Attention module with to_q/to_k/to_v/add_*_proj, norms, to_out, to_add_out, and .heads
+        hidden_states: torch.FloatTensor,                 # (B, S_img, D_model)  image stream
+        encoder_hidden_states: torch.FloatTensor = None,  # (B, S_txt, D_model)  text stream
+        encoder_hidden_states_mask: torch.FloatTensor = None,  # unused in FA3 path
+        attention_mask: Optional[torch.FloatTensor] = None,    # unused in FA3 path
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # (img_freqs, txt_freqs)
+    ) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
+        if encoder_hidden_states is None:
+            raise ValueError("QwenDoubleStreamAttnProcessorFA3 requires encoder_hidden_states (text stream).")
+        if attention_mask is not None:
+            raise NotImplementedError("attention_mask is not supported in this FA3 implementation.")
+
+        B, S_img, _ = hidden_states.shape
+        S_txt = encoder_hidden_states.shape[1]
+
+        # ---- QKV projections (image/sample stream) ----
+        img_q = attn.to_q(hidden_states)   # (B, S_img, D)
+        img_k = attn.to_k(hidden_states)
+        img_v = attn.to_v(hidden_states)
+
+        # ---- QKV projections (text/context stream) ----
+        txt_q = attn.add_q_proj(encoder_hidden_states)  # (B, S_txt, D)
+        txt_k = attn.add_k_proj(encoder_hidden_states)
+        txt_v = attn.add_v_proj(encoder_hidden_states)
+
+        # ---- Reshape to (B, S, H, D_h) ----
+        H = attn.heads
+        img_q = img_q.unflatten(-1, (H, -1))
+        img_k = img_k.unflatten(-1, (H, -1))
+        img_v = img_v.unflatten(-1, (H, -1))
+
+        txt_q = txt_q.unflatten(-1, (H, -1))
+        txt_k = txt_k.unflatten(-1, (H, -1))
+        txt_v = txt_v.unflatten(-1, (H, -1))
+
+        # ---- Q/K normalization (per your module contract) ----
+        if getattr(attn, "norm_q", None) is not None:
+            img_q = attn.norm_q(img_q)
+        if getattr(attn, "norm_k", None) is not None:
+            img_k = attn.norm_k(img_k)
+        if getattr(attn, "norm_added_q", None) is not None:
+            txt_q = attn.norm_added_q(txt_q)
+        if getattr(attn, "norm_added_k", None) is not None:
+            txt_k = attn.norm_added_k(txt_k)
+
+        # ---- RoPE (Qwen variant) ----
+        if image_rotary_emb is not None:
+            img_freqs, txt_freqs = image_rotary_emb
+            img_q = apply_rotary_emb_qwen(img_q, img_freqs, use_real=False)
+            img_k = apply_rotary_emb_qwen(img_k, img_freqs, use_real=False)
+            txt_q = apply_rotary_emb_qwen(txt_q, txt_freqs, use_real=False)
+            txt_k = apply_rotary_emb_qwen(txt_k, txt_freqs, use_real=False)
+
+        # ---- Joint attention over [text, image] along sequence axis ----
+        q = torch.cat([txt_q, img_q], dim=1)  # (B, S_total, H, D_h)
+        k = torch.cat([txt_k, img_k], dim=1)
+        v = torch.cat([txt_v, img_v], dim=1)
+
+        # SDPA wants (B, H, S, D_h); route through cuDNN's FA3 path on Hopper.
+        with sdpa_kernel(SDPBackend.CUDNN_ATTENTION):
+            out = F.scaled_dot_product_attention(
+                q.transpose(1, 2),
+                k.transpose(1, 2),
+                v.transpose(1, 2),
+                is_causal=False,
+            ).transpose(1, 2)  # back to (B, S_total, H, D_h)
+
+        # ---- Back to (B, S, D_model) ----
+        out = out.flatten(2, 3).to(q.dtype)
+
+        # Split back to text / image segments
+        txt_attn_out = out[:, :S_txt, :]
+        img_attn_out = out[:, S_txt:, :]
+
+        # ---- Output projections ----
+        img_attn_out = attn.to_out[0](img_attn_out)
+        if len(attn.to_out) > 1:
+            img_attn_out = attn.to_out[1](img_attn_out)  # dropout if present
+
+        txt_attn_out = attn.to_add_out(txt_attn_out)
+
+        return img_attn_out, txt_attn_out

qwenimage/transformer_qwenimage.py

@@ -1,642 +1,642 @@
-# Copyright 2025 Qwen-Image Team, The HuggingFace Team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-import functools
-import math
-from typing import Any, Dict, List, Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
-from diffusers.utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
-from diffusers.utils.torch_utils import maybe_allow_in_graph
-from diffusers.models.attention import FeedForward, AttentionMixin
-from diffusers.models.attention_dispatch import dispatch_attention_fn
-from diffusers.models.attention_processor import Attention
-from diffusers.models.cache_utils import CacheMixin
-from diffusers.models.embeddings import TimestepEmbedding, Timesteps
-from diffusers.models.modeling_outputs import Transformer2DModelOutput
-from diffusers.models.modeling_utils import ModelMixin
-from diffusers.models.normalization import AdaLayerNormContinuous, RMSNorm
-
-
-logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
-
-
-def get_timestep_embedding(
-    timesteps: torch.Tensor,
-    embedding_dim: int,
-    flip_sin_to_cos: bool = False,
-    downscale_freq_shift: float = 1,
-    scale: float = 1,
-    max_period: int = 10000,
-) -> torch.Tensor:
-    """
-    This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
-
-    Args
-        timesteps (torch.Tensor):
-            a 1-D Tensor of N indices, one per batch element. These may be fractional.
-        embedding_dim (int):
-            the dimension of the output.
-        flip_sin_to_cos (bool):
-            Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
-        downscale_freq_shift (float):
-            Controls the delta between frequencies between dimensions
-        scale (float):
-            Scaling factor applied to the embeddings.
-        max_period (int):
-            Controls the maximum frequency of the embeddings
-    Returns
-        torch.Tensor: an [N x dim] Tensor of positional embeddings.
-    """
-    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
-
-    half_dim = embedding_dim // 2
-    exponent = -math.log(max_period) * torch.arange(
-        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
-    )
-    exponent = exponent / (half_dim - downscale_freq_shift)
-
-    emb = torch.exp(exponent).to(timesteps.dtype)
-    emb = timesteps[:, None].float() * emb[None, :]
-
-    # scale embeddings
-    emb = scale * emb
-
-    # concat sine and cosine embeddings
-    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
-
-    # flip sine and cosine embeddings
-    if flip_sin_to_cos:
-        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
-
-    # zero pad
-    if embedding_dim % 2 == 1:
-        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
-    return emb
-
-
-def apply_rotary_emb_qwen(
-    x: torch.Tensor,
-    freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
-    use_real: bool = True,
-    use_real_unbind_dim: int = -1,
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    """
-    Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
-    to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
-    reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
-    tensors contain rotary embeddings and are returned as real tensors.
-
-    Args:
-        x (`torch.Tensor`):
-            Query or key tensor to apply rotary embeddings. [B, S, H, D] xk (torch.Tensor): Key tensor to apply
-        freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
-
-    Returns:
-        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
-    """
-    if use_real:
-        cos, sin = freqs_cis  # [S, D]
-        cos = cos[None, None]
-        sin = sin[None, None]
-        cos, sin = cos.to(x.device), sin.to(x.device)
-
-        if use_real_unbind_dim == -1:
-            # Used for flux, cogvideox, hunyuan-dit
-            x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
-            x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
-        elif use_real_unbind_dim == -2:
-            # Used for Stable Audio, OmniGen, CogView4 and Cosmos
-            x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2)  # [B, S, H, D//2]
-            x_rotated = torch.cat([-x_imag, x_real], dim=-1)
-        else:
-            raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
-
-        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
-
-        return out
-    else:
-        x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
-        freqs_cis = freqs_cis.unsqueeze(1)
-        x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
-
-        return x_out.type_as(x)
-
-
-class QwenTimestepProjEmbeddings(nn.Module):
-    def __init__(self, embedding_dim):
-        super().__init__()
-
-        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000)
-        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
-
-    def forward(self, timestep, hidden_states):
-        timesteps_proj = self.time_proj(timestep)
-        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_states.dtype))  # (N, D)
-
-        conditioning = timesteps_emb
-
-        return conditioning
-
-
-class QwenEmbedRope(nn.Module):
-    def __init__(self, theta: int, axes_dim: List[int], scale_rope=False):
-        super().__init__()
-        self.theta = theta
-        self.axes_dim = axes_dim
-        pos_index = torch.arange(4096)
-        neg_index = torch.arange(4096).flip(0) * -1 - 1
-        self.pos_freqs = torch.cat(
-            [
-                self.rope_params(pos_index, self.axes_dim[0], self.theta),
-                self.rope_params(pos_index, self.axes_dim[1], self.theta),
-                self.rope_params(pos_index, self.axes_dim[2], self.theta),
-            ],
-            dim=1,
-        )
-        self.neg_freqs = torch.cat(
-            [
-                self.rope_params(neg_index, self.axes_dim[0], self.theta),
-                self.rope_params(neg_index, self.axes_dim[1], self.theta),
-                self.rope_params(neg_index, self.axes_dim[2], self.theta),
-            ],
-            dim=1,
-        )
-        self.rope_cache = {}
-
-        # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
-        self.scale_rope = scale_rope
-
-    def rope_params(self, index, dim, theta=10000):
-        """
-        Args:
-            index: [0, 1, 2, 3] 1D Tensor representing the position index of the token
-        """
-        assert dim % 2 == 0
-        freqs = torch.outer(index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim)))
-        freqs = torch.polar(torch.ones_like(freqs), freqs)
-        return freqs
-
-    def forward(self, video_fhw, txt_seq_lens, device):
-        """
-        Args: video_fhw: [frame, height, width] a list of 3 integers representing the shape of the video Args:
-        txt_length: [bs] a list of 1 integers representing the length of the text
-        """
-        if self.pos_freqs.device != device:
-            self.pos_freqs = self.pos_freqs.to(device)
-            self.neg_freqs = self.neg_freqs.to(device)
-
-        if isinstance(video_fhw, list):
-            video_fhw = video_fhw[0]
-        if not isinstance(video_fhw, list):
-            video_fhw = [video_fhw]
-
-        vid_freqs = []
-        max_vid_index = 0
-        for idx, fhw in enumerate(video_fhw):
-            frame, height, width = fhw
-            rope_key = f"{idx}_{height}_{width}"
-
-            if not torch.compiler.is_compiling():
-                if rope_key not in self.rope_cache:
-                    self.rope_cache[rope_key] = self._compute_video_freqs(frame, height, width, idx)
-                video_freq = self.rope_cache[rope_key]
-            else:
-                video_freq = self._compute_video_freqs(frame, height, width, idx)
-            video_freq = video_freq.to(device)
-            vid_freqs.append(video_freq)
-
-            if self.scale_rope:
-                max_vid_index = max(height // 2, width // 2, max_vid_index)
-            else:
-                max_vid_index = max(height, width, max_vid_index)
-
-        max_len = max(txt_seq_lens)
-        txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...]
-        vid_freqs = torch.cat(vid_freqs, dim=0)
-
-        return vid_freqs, txt_freqs
-
-    @functools.lru_cache(maxsize=None)
-    def _compute_video_freqs(self, frame, height, width, idx=0):
-        seq_lens = frame * height * width
-        freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
-        freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)
-
-        freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1)
-        if self.scale_rope:
-            freqs_height = torch.cat([freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], dim=0)
-            freqs_height = freqs_height.view(1, height, 1, -1).expand(frame, height, width, -1)
-            freqs_width = torch.cat([freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], dim=0)
-            freqs_width = freqs_width.view(1, 1, width, -1).expand(frame, height, width, -1)
-        else:
-            freqs_height = freqs_pos[1][:height].view(1, height, 1, -1).expand(frame, height, width, -1)
-            freqs_width = freqs_pos[2][:width].view(1, 1, width, -1).expand(frame, height, width, -1)
-
-        freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1)
-        return freqs.clone().contiguous()
-
-
-class QwenDoubleStreamAttnProcessor2_0:
-    """
-    Attention processor for Qwen double-stream architecture, matching DoubleStreamLayerMegatron logic. This processor
-    implements joint attention computation where text and image streams are processed together.
-    """
-
-    _attention_backend = None
-
-    def __init__(self):
-        if not hasattr(F, "scaled_dot_product_attention"):
-            raise ImportError(
-                "QwenDoubleStreamAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
-            )
-
-    def __call__(
-        self,
-        attn: Attention,
-        hidden_states: torch.FloatTensor,  # Image stream
-        encoder_hidden_states: torch.FloatTensor = None,  # Text stream
-        encoder_hidden_states_mask: torch.FloatTensor = None,
-        attention_mask: Optional[torch.FloatTensor] = None,
-        image_rotary_emb: Optional[torch.Tensor] = None,
-    ) -> torch.FloatTensor:
-        if encoder_hidden_states is None:
-            raise ValueError("QwenDoubleStreamAttnProcessor2_0 requires encoder_hidden_states (text stream)")
-
-        seq_txt = encoder_hidden_states.shape[1]
-
-        # Compute QKV for image stream (sample projections)
-        img_query = attn.to_q(hidden_states)
-        img_key = attn.to_k(hidden_states)
-        img_value = attn.to_v(hidden_states)
-
-        # Compute QKV for text stream (context projections)
-        txt_query = attn.add_q_proj(encoder_hidden_states)
-        txt_key = attn.add_k_proj(encoder_hidden_states)
-        txt_value = attn.add_v_proj(encoder_hidden_states)
-
-        # Reshape for multi-head attention
-        img_query = img_query.unflatten(-1, (attn.heads, -1))
-        img_key = img_key.unflatten(-1, (attn.heads, -1))
-        img_value = img_value.unflatten(-1, (attn.heads, -1))
-
-        txt_query = txt_query.unflatten(-1, (attn.heads, -1))
-        txt_key = txt_key.unflatten(-1, (attn.heads, -1))
-        txt_value = txt_value.unflatten(-1, (attn.heads, -1))
-
-        # Apply QK normalization
-        if attn.norm_q is not None:
-            img_query = attn.norm_q(img_query)
-        if attn.norm_k is not None:
-            img_key = attn.norm_k(img_key)
-        if attn.norm_added_q is not None:
-            txt_query = attn.norm_added_q(txt_query)
-        if attn.norm_added_k is not None:
-            txt_key = attn.norm_added_k(txt_key)
-
-        # Apply RoPE
-        if image_rotary_emb is not None:
-            img_freqs, txt_freqs = image_rotary_emb
-            img_query = apply_rotary_emb_qwen(img_query, img_freqs, use_real=False)
-            img_key = apply_rotary_emb_qwen(img_key, img_freqs, use_real=False)
-            txt_query = apply_rotary_emb_qwen(txt_query, txt_freqs, use_real=False)
-            txt_key = apply_rotary_emb_qwen(txt_key, txt_freqs, use_real=False)
-
-        # Concatenate for joint attention
-        # Order: [text, image]
-        joint_query = torch.cat([txt_query, img_query], dim=1)
-        joint_key = torch.cat([txt_key, img_key], dim=1)
-        joint_value = torch.cat([txt_value, img_value], dim=1)
-
-        # Compute joint attention
-        joint_hidden_states = dispatch_attention_fn(
-            joint_query,
-            joint_key,
-            joint_value,
-            attn_mask=attention_mask,
-            dropout_p=0.0,
-            is_causal=False,
-            backend=self._attention_backend,
-        )
-
-        # Reshape back
-        joint_hidden_states = joint_hidden_states.flatten(2, 3)
-        joint_hidden_states = joint_hidden_states.to(joint_query.dtype)
-
-        # Split attention outputs back
-        txt_attn_output = joint_hidden_states[:, :seq_txt, :]  # Text part
-        img_attn_output = joint_hidden_states[:, seq_txt:, :]  # Image part
-
-        # Apply output projections
-        img_attn_output = attn.to_out[0](img_attn_output)
-        if len(attn.to_out) > 1:
-            img_attn_output = attn.to_out[1](img_attn_output)  # dropout
-
-        txt_attn_output = attn.to_add_out(txt_attn_output)
-
-        return img_attn_output, txt_attn_output
-
-
-@maybe_allow_in_graph
-class QwenImageTransformerBlock(nn.Module):
-    def __init__(
-        self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6
-    ):
-        super().__init__()
-
-        self.dim = dim
-        self.num_attention_heads = num_attention_heads
-        self.attention_head_dim = attention_head_dim
-
-        # Image processing modules
-        self.img_mod = nn.Sequential(
-            nn.SiLU(),
-            nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
-        )
-        self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.attn = Attention(
-            query_dim=dim,
-            cross_attention_dim=None,  # Enable cross attention for joint computation
-            added_kv_proj_dim=dim,  # Enable added KV projections for text stream
-            dim_head=attention_head_dim,
-            heads=num_attention_heads,
-            out_dim=dim,
-            context_pre_only=False,
-            bias=True,
-            processor=QwenDoubleStreamAttnProcessor2_0(),
-            qk_norm=qk_norm,
-            eps=eps,
-        )
-        self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.img_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
-
-        # Text processing modules
-        self.txt_mod = nn.Sequential(
-            nn.SiLU(),
-            nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
-        )
-        self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
-        # Text doesn't need separate attention - it's handled by img_attn joint computation
-        self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
-        self.txt_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
-
-    def _modulate(self, x, mod_params):
-        """Apply modulation to input tensor"""
-        shift, scale, gate = mod_params.chunk(3, dim=-1)
-        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1), gate.unsqueeze(1)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        encoder_hidden_states_mask: torch.Tensor,
-        temb: torch.Tensor,
-        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
-        # Get modulation parameters for both streams
-        img_mod_params = self.img_mod(temb)  # [B, 6*dim]
-        txt_mod_params = self.txt_mod(temb)  # [B, 6*dim]
-
-        # Split modulation parameters for norm1 and norm2
-        img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
-        txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
-
-        # Process image stream - norm1 + modulation
-        img_normed = self.img_norm1(hidden_states)
-        img_modulated, img_gate1 = self._modulate(img_normed, img_mod1)
-
-        # Process text stream - norm1 + modulation
-        txt_normed = self.txt_norm1(encoder_hidden_states)
-        txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1)
-
-        # Use QwenAttnProcessor2_0 for joint attention computation
-        # This directly implements the DoubleStreamLayerMegatron logic:
-        # 1. Computes QKV for both streams
-        # 2. Applies QK normalization and RoPE
-        # 3. Concatenates and runs joint attention
-        # 4. Splits results back to separate streams
-        joint_attention_kwargs = joint_attention_kwargs or {}
-        attn_output = self.attn(
-            hidden_states=img_modulated,  # Image stream (will be processed as "sample")
-            encoder_hidden_states=txt_modulated,  # Text stream (will be processed as "context")
-            encoder_hidden_states_mask=encoder_hidden_states_mask,
-            image_rotary_emb=image_rotary_emb,
-            **joint_attention_kwargs,
-        )
-
-        # QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided
-        img_attn_output, txt_attn_output = attn_output
-
-        # Apply attention gates and add residual (like in Megatron)
-        hidden_states = hidden_states + img_gate1 * img_attn_output
-        encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
-
-        # Process image stream - norm2 + MLP
-        img_normed2 = self.img_norm2(hidden_states)
-        img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2)
-        img_mlp_output = self.img_mlp(img_modulated2)
-        hidden_states = hidden_states + img_gate2 * img_mlp_output
-
-        # Process text stream - norm2 + MLP
-        txt_normed2 = self.txt_norm2(encoder_hidden_states)
-        txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2)
-        txt_mlp_output = self.txt_mlp(txt_modulated2)
-        encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output
-
-        # Clip to prevent overflow for fp16
-        if encoder_hidden_states.dtype == torch.float16:
-            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
-        if hidden_states.dtype == torch.float16:
-            hidden_states = hidden_states.clip(-65504, 65504)
-
-        return encoder_hidden_states, hidden_states
-
-
-class QwenImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin, AttentionMixin):
-    """
-    The Transformer model introduced in Qwen.
-
-    Args:
-        patch_size (`int`, defaults to `2`):
-            Patch size to turn the input data into small patches.
-        in_channels (`int`, defaults to `64`):
-            The number of channels in the input.
-        out_channels (`int`, *optional*, defaults to `None`):
-            The number of channels in the output. If not specified, it defaults to `in_channels`.
-        num_layers (`int`, defaults to `60`):
-            The number of layers of dual stream DiT blocks to use.
-        attention_head_dim (`int`, defaults to `128`):
-            The number of dimensions to use for each attention head.
-        num_attention_heads (`int`, defaults to `24`):
-            The number of attention heads to use.
-        joint_attention_dim (`int`, defaults to `3584`):
-            The number of dimensions to use for the joint attention (embedding/channel dimension of
-            `encoder_hidden_states`).
-        guidance_embeds (`bool`, defaults to `False`):
-            Whether to use guidance embeddings for guidance-distilled variant of the model.
-        axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`):
-            The dimensions to use for the rotary positional embeddings.
-    """
-
-    _supports_gradient_checkpointing = True
-    _no_split_modules = ["QwenImageTransformerBlock"]
-    _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
-    _repeated_blocks = ["QwenImageTransformerBlock"]
-
-    @register_to_config
-    def __init__(
-        self,
-        patch_size: int = 2,
-        in_channels: int = 64,
-        out_channels: Optional[int] = 16,
-        num_layers: int = 60,
-        attention_head_dim: int = 128,
-        num_attention_heads: int = 24,
-        joint_attention_dim: int = 3584,
-        guidance_embeds: bool = False,  # TODO: this should probably be removed
-        axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
-    ):
-        super().__init__()
-        self.out_channels = out_channels or in_channels
-        self.inner_dim = num_attention_heads * attention_head_dim
-
-        self.pos_embed = QwenEmbedRope(theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True)
-
-        self.time_text_embed = QwenTimestepProjEmbeddings(embedding_dim=self.inner_dim)
-
-        self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
-
-        self.img_in = nn.Linear(in_channels, self.inner_dim)
-        self.txt_in = nn.Linear(joint_attention_dim, self.inner_dim)
-
-        self.transformer_blocks = nn.ModuleList(
-            [
-                QwenImageTransformerBlock(
-                    dim=self.inner_dim,
-                    num_attention_heads=num_attention_heads,
-                    attention_head_dim=attention_head_dim,
-                )
-                for _ in range(num_layers)
-            ]
-        )
-
-        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
-        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
-
-        self.gradient_checkpointing = False
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor = None,
-        encoder_hidden_states_mask: torch.Tensor = None,
-        timestep: torch.LongTensor = None,
-        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
-        guidance: torch.Tensor = None,  # TODO: this should probably be removed
-        attention_kwargs: Optional[Dict[str, Any]] = None,
-        return_dict: bool = True,
-    ) -> Union[torch.Tensor, Transformer2DModelOutput]:
-        """
-        The [`QwenTransformer2DModel`] forward method.
-
-        Args:
-            hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
-                Input `hidden_states`.
-            encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
-                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
-            encoder_hidden_states_mask (`torch.Tensor` of shape `(batch_size, text_sequence_length)`):
-                Mask of the input conditions.
-            timestep ( `torch.LongTensor`):
-                Used to indicate denoising step.
-            attention_kwargs (`dict`, *optional*):
-                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
-                `self.processor` in
-                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
-            return_dict (`bool`, *optional*, defaults to `True`):
-                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
-                tuple.
-
-        Returns:
-            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
-            `tuple` where the first element is the sample tensor.
-        """
-        if attention_kwargs is not None:
-            attention_kwargs = attention_kwargs.copy()
-            lora_scale = attention_kwargs.pop("scale", 1.0)
-        else:
-            lora_scale = 1.0
-
-        if USE_PEFT_BACKEND:
-            # weight the lora layers by setting `lora_scale` for each PEFT layer
-            scale_lora_layers(self, lora_scale)
-        else:
-            if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
-                logger.warning(
-                    "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
-                )
-
-        hidden_states = self.img_in(hidden_states)
-
-        timestep = timestep.to(hidden_states.dtype)
-        encoder_hidden_states = self.txt_norm(encoder_hidden_states)
-        encoder_hidden_states = self.txt_in(encoder_hidden_states)
-
-        if guidance is not None:
-            guidance = guidance.to(hidden_states.dtype) * 1000
-
-        temb = (
-            self.time_text_embed(timestep, hidden_states)
-            if guidance is None
-            else self.time_text_embed(timestep, guidance, hidden_states)
-        )
-
-        for index_block, block in enumerate(self.transformer_blocks):
-            if torch.is_grad_enabled() and self.gradient_checkpointing:
-                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
-                    block,
-                    hidden_states,
-                    encoder_hidden_states,
-                    encoder_hidden_states_mask,
-                    temb,
-                    image_rotary_emb,
-                )
-
-            else:
-                encoder_hidden_states, hidden_states = block(
-                    hidden_states=hidden_states,
-                    encoder_hidden_states=encoder_hidden_states,
-                    encoder_hidden_states_mask=encoder_hidden_states_mask,
-                    temb=temb,
-                    image_rotary_emb=image_rotary_emb,
-                    joint_attention_kwargs=attention_kwargs,
-                )
-
-        # Use only the image part (hidden_states) from the dual-stream blocks
-        hidden_states = self.norm_out(hidden_states, temb)
-        output = self.proj_out(hidden_states)
-
-        if USE_PEFT_BACKEND:
-            # remove `lora_scale` from each PEFT layer
-            unscale_lora_layers(self, lora_scale)
-
-        if not return_dict:
-            return (output,)
-
-        return Transformer2DModelOutput(sample=output)
+# Copyright 2025 Qwen-Image Team, The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import functools
+import math
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
+from diffusers.utils import USE_PEFT_BACKEND, logging, scale_lora_layers, unscale_lora_layers
+from diffusers.utils.torch_utils import maybe_allow_in_graph
+from diffusers.models.attention import FeedForward, AttentionMixin
+from diffusers.models.attention_dispatch import dispatch_attention_fn
+from diffusers.models.attention_processor import Attention
+from diffusers.models.cache_utils import CacheMixin
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+from diffusers.models.modeling_outputs import Transformer2DModelOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.models.normalization import AdaLayerNormContinuous, RMSNorm
+
+
+logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
+
+
+def get_timestep_embedding(
+    timesteps: torch.Tensor,
+    embedding_dim: int,
+    flip_sin_to_cos: bool = False,
+    downscale_freq_shift: float = 1,
+    scale: float = 1,
+    max_period: int = 10000,
+) -> torch.Tensor:
+    """
+    This matches the implementation in Denoising Diffusion Probabilistic Models: Create sinusoidal timestep embeddings.
+
+    Args
+        timesteps (torch.Tensor):
+            a 1-D Tensor of N indices, one per batch element. These may be fractional.
+        embedding_dim (int):
+            the dimension of the output.
+        flip_sin_to_cos (bool):
+            Whether the embedding order should be `cos, sin` (if True) or `sin, cos` (if False)
+        downscale_freq_shift (float):
+            Controls the delta between frequencies between dimensions
+        scale (float):
+            Scaling factor applied to the embeddings.
+        max_period (int):
+            Controls the maximum frequency of the embeddings
+    Returns
+        torch.Tensor: an [N x dim] Tensor of positional embeddings.
+    """
+    assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+
+    half_dim = embedding_dim // 2
+    exponent = -math.log(max_period) * torch.arange(
+        start=0, end=half_dim, dtype=torch.float32, device=timesteps.device
+    )
+    exponent = exponent / (half_dim - downscale_freq_shift)
+
+    emb = torch.exp(exponent).to(timesteps.dtype)
+    emb = timesteps[:, None].float() * emb[None, :]
+
+    # scale embeddings
+    emb = scale * emb
+
+    # concat sine and cosine embeddings
+    emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+
+    # flip sine and cosine embeddings
+    if flip_sin_to_cos:
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
+
+    # zero pad
+    if embedding_dim % 2 == 1:
+        emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+    return emb
+
+
+def apply_rotary_emb_qwen(
+    x: torch.Tensor,
+    freqs_cis: Union[torch.Tensor, Tuple[torch.Tensor]],
+    use_real: bool = True,
+    use_real_unbind_dim: int = -1,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensor. This function applies rotary embeddings
+    to the given query or key 'x' tensors using the provided frequency tensor 'freqs_cis'. The input tensors are
+    reshaped as complex numbers, and the frequency tensor is reshaped for broadcasting compatibility. The resulting
+    tensors contain rotary embeddings and are returned as real tensors.
+
+    Args:
+        x (`torch.Tensor`):
+            Query or key tensor to apply rotary embeddings. [B, S, H, D] xk (torch.Tensor): Key tensor to apply
+        freqs_cis (`Tuple[torch.Tensor]`): Precomputed frequency tensor for complex exponentials. ([S, D], [S, D],)
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+    """
+    if use_real:
+        cos, sin = freqs_cis  # [S, D]
+        cos = cos[None, None]
+        sin = sin[None, None]
+        cos, sin = cos.to(x.device), sin.to(x.device)
+
+        if use_real_unbind_dim == -1:
+            # Used for flux, cogvideox, hunyuan-dit
+            x_real, x_imag = x.reshape(*x.shape[:-1], -1, 2).unbind(-1)  # [B, S, H, D//2]
+            x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(3)
+        elif use_real_unbind_dim == -2:
+            # Used for Stable Audio, OmniGen, CogView4 and Cosmos
+            x_real, x_imag = x.reshape(*x.shape[:-1], 2, -1).unbind(-2)  # [B, S, H, D//2]
+            x_rotated = torch.cat([-x_imag, x_real], dim=-1)
+        else:
+            raise ValueError(f"`use_real_unbind_dim={use_real_unbind_dim}` but should be -1 or -2.")
+
+        out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+
+        return out
+    else:
+        x_rotated = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        freqs_cis = freqs_cis.unsqueeze(1)
+        x_out = torch.view_as_real(x_rotated * freqs_cis).flatten(3)
+
+        return x_out.type_as(x)
+
+
+class QwenTimestepProjEmbeddings(nn.Module):
+    def __init__(self, embedding_dim):
+        super().__init__()
+
+        self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000)
+        self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)
+
+    def forward(self, timestep, hidden_states):
+        timesteps_proj = self.time_proj(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_states.dtype))  # (N, D)
+
+        conditioning = timesteps_emb
+
+        return conditioning
+
+
+class QwenEmbedRope(nn.Module):
+    def __init__(self, theta: int, axes_dim: List[int], scale_rope=False):
+        super().__init__()
+        self.theta = theta
+        self.axes_dim = axes_dim
+        pos_index = torch.arange(4096)
+        neg_index = torch.arange(4096).flip(0) * -1 - 1
+        self.pos_freqs = torch.cat(
+            [
+                self.rope_params(pos_index, self.axes_dim[0], self.theta),
+                self.rope_params(pos_index, self.axes_dim[1], self.theta),
+                self.rope_params(pos_index, self.axes_dim[2], self.theta),
+            ],
+            dim=1,
+        )
+        self.neg_freqs = torch.cat(
+            [
+                self.rope_params(neg_index, self.axes_dim[0], self.theta),
+                self.rope_params(neg_index, self.axes_dim[1], self.theta),
+                self.rope_params(neg_index, self.axes_dim[2], self.theta),
+            ],
+            dim=1,
+        )
+        self.rope_cache = {}
+
+        # DO NOT USING REGISTER BUFFER HERE, IT WILL CAUSE COMPLEX NUMBERS LOSE ITS IMAGINARY PART
+        self.scale_rope = scale_rope
+
+    def rope_params(self, index, dim, theta=10000):
+        """
+        Args:
+            index: [0, 1, 2, 3] 1D Tensor representing the position index of the token
+        """
+        assert dim % 2 == 0
+        freqs = torch.outer(index, 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float32).div(dim)))
+        freqs = torch.polar(torch.ones_like(freqs), freqs)
+        return freqs
+
+    def forward(self, video_fhw, txt_seq_lens, device):
+        """
+        Args: video_fhw: [frame, height, width] a list of 3 integers representing the shape of the video Args:
+        txt_length: [bs] a list of 1 integers representing the length of the text
+        """
+        if self.pos_freqs.device != device:
+            self.pos_freqs = self.pos_freqs.to(device)
+            self.neg_freqs = self.neg_freqs.to(device)
+
+        if isinstance(video_fhw, list):
+            video_fhw = video_fhw[0]
+        if not isinstance(video_fhw, list):
+            video_fhw = [video_fhw]
+
+        vid_freqs = []
+        max_vid_index = 0
+        for idx, fhw in enumerate(video_fhw):
+            frame, height, width = fhw
+            rope_key = f"{idx}_{height}_{width}"
+
+            if not torch.compiler.is_compiling():
+                if rope_key not in self.rope_cache:
+                    self.rope_cache[rope_key] = self._compute_video_freqs(frame, height, width, idx)
+                video_freq = self.rope_cache[rope_key]
+            else:
+                video_freq = self._compute_video_freqs(frame, height, width, idx)
+            video_freq = video_freq.to(device)
+            vid_freqs.append(video_freq)
+
+            if self.scale_rope:
+                max_vid_index = max(height // 2, width // 2, max_vid_index)
+            else:
+                max_vid_index = max(height, width, max_vid_index)
+
+        max_len = max(txt_seq_lens)
+        txt_freqs = self.pos_freqs[max_vid_index : max_vid_index + max_len, ...]
+        vid_freqs = torch.cat(vid_freqs, dim=0)
+
+        return vid_freqs, txt_freqs
+
+    @functools.lru_cache(maxsize=None)
+    def _compute_video_freqs(self, frame, height, width, idx=0):
+        seq_lens = frame * height * width
+        freqs_pos = self.pos_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+        freqs_neg = self.neg_freqs.split([x // 2 for x in self.axes_dim], dim=1)
+
+        freqs_frame = freqs_pos[0][idx : idx + frame].view(frame, 1, 1, -1).expand(frame, height, width, -1)
+        if self.scale_rope:
+            freqs_height = torch.cat([freqs_neg[1][-(height - height // 2) :], freqs_pos[1][: height // 2]], dim=0)
+            freqs_height = freqs_height.view(1, height, 1, -1).expand(frame, height, width, -1)
+            freqs_width = torch.cat([freqs_neg[2][-(width - width // 2) :], freqs_pos[2][: width // 2]], dim=0)
+            freqs_width = freqs_width.view(1, 1, width, -1).expand(frame, height, width, -1)
+        else:
+            freqs_height = freqs_pos[1][:height].view(1, height, 1, -1).expand(frame, height, width, -1)
+            freqs_width = freqs_pos[2][:width].view(1, 1, width, -1).expand(frame, height, width, -1)
+
+        freqs = torch.cat([freqs_frame, freqs_height, freqs_width], dim=-1).reshape(seq_lens, -1)
+        return freqs.clone().contiguous()
+
+
+class QwenDoubleStreamAttnProcessor2_0:
+    """
+    Attention processor for Qwen double-stream architecture, matching DoubleStreamLayerMegatron logic. This processor
+    implements joint attention computation where text and image streams are processed together.
+    """
+
+    _attention_backend = None
+
+    def __init__(self):
+        if not hasattr(F, "scaled_dot_product_attention"):
+            raise ImportError(
+                "QwenDoubleStreamAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+            )
+
+    def __call__(
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,  # Image stream
+        encoder_hidden_states: torch.FloatTensor = None,  # Text stream
+        encoder_hidden_states_mask: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+    ) -> torch.FloatTensor:
+        if encoder_hidden_states is None:
+            raise ValueError("QwenDoubleStreamAttnProcessor2_0 requires encoder_hidden_states (text stream)")
+
+        seq_txt = encoder_hidden_states.shape[1]
+
+        # Compute QKV for image stream (sample projections)
+        img_query = attn.to_q(hidden_states)
+        img_key = attn.to_k(hidden_states)
+        img_value = attn.to_v(hidden_states)
+
+        # Compute QKV for text stream (context projections)
+        txt_query = attn.add_q_proj(encoder_hidden_states)
+        txt_key = attn.add_k_proj(encoder_hidden_states)
+        txt_value = attn.add_v_proj(encoder_hidden_states)
+
+        # Reshape for multi-head attention
+        img_query = img_query.unflatten(-1, (attn.heads, -1))
+        img_key = img_key.unflatten(-1, (attn.heads, -1))
+        img_value = img_value.unflatten(-1, (attn.heads, -1))
+
+        txt_query = txt_query.unflatten(-1, (attn.heads, -1))
+        txt_key = txt_key.unflatten(-1, (attn.heads, -1))
+        txt_value = txt_value.unflatten(-1, (attn.heads, -1))
+
+        # Apply QK normalization
+        if attn.norm_q is not None:
+            img_query = attn.norm_q(img_query)
+        if attn.norm_k is not None:
+            img_key = attn.norm_k(img_key)
+        if attn.norm_added_q is not None:
+            txt_query = attn.norm_added_q(txt_query)
+        if attn.norm_added_k is not None:
+            txt_key = attn.norm_added_k(txt_key)
+
+        # Apply RoPE
+        if image_rotary_emb is not None:
+            img_freqs, txt_freqs = image_rotary_emb
+            img_query = apply_rotary_emb_qwen(img_query, img_freqs, use_real=False)
+            img_key = apply_rotary_emb_qwen(img_key, img_freqs, use_real=False)
+            txt_query = apply_rotary_emb_qwen(txt_query, txt_freqs, use_real=False)
+            txt_key = apply_rotary_emb_qwen(txt_key, txt_freqs, use_real=False)
+
+        # Concatenate for joint attention
+        # Order: [text, image]
+        joint_query = torch.cat([txt_query, img_query], dim=1)
+        joint_key = torch.cat([txt_key, img_key], dim=1)
+        joint_value = torch.cat([txt_value, img_value], dim=1)
+
+        # Compute joint attention
+        joint_hidden_states = dispatch_attention_fn(
+            joint_query,
+            joint_key,
+            joint_value,
+            attn_mask=attention_mask,
+            dropout_p=0.0,
+            is_causal=False,
+            backend=self._attention_backend,
+        )
+
+        # Reshape back
+        joint_hidden_states = joint_hidden_states.flatten(2, 3)
+        joint_hidden_states = joint_hidden_states.to(joint_query.dtype)
+
+        # Split attention outputs back
+        txt_attn_output = joint_hidden_states[:, :seq_txt, :]  # Text part
+        img_attn_output = joint_hidden_states[:, seq_txt:, :]  # Image part
+
+        # Apply output projections
+        img_attn_output = attn.to_out[0](img_attn_output)
+        if len(attn.to_out) > 1:
+            img_attn_output = attn.to_out[1](img_attn_output)  # dropout
+
+        txt_attn_output = attn.to_add_out(txt_attn_output)
+
+        return img_attn_output, txt_attn_output
+
+
+@maybe_allow_in_graph
+class QwenImageTransformerBlock(nn.Module):
+    def __init__(
+        self, dim: int, num_attention_heads: int, attention_head_dim: int, qk_norm: str = "rms_norm", eps: float = 1e-6
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.num_attention_heads = num_attention_heads
+        self.attention_head_dim = attention_head_dim
+
+        # Image processing modules
+        self.img_mod = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
+        )
+        self.img_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.attn = Attention(
+            query_dim=dim,
+            cross_attention_dim=None,  # Enable cross attention for joint computation
+            added_kv_proj_dim=dim,  # Enable added KV projections for text stream
+            dim_head=attention_head_dim,
+            heads=num_attention_heads,
+            out_dim=dim,
+            context_pre_only=False,
+            bias=True,
+            processor=QwenDoubleStreamAttnProcessor2_0(),
+            qk_norm=qk_norm,
+            eps=eps,
+        )
+        self.img_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.img_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+        # Text processing modules
+        self.txt_mod = nn.Sequential(
+            nn.SiLU(),
+            nn.Linear(dim, 6 * dim, bias=True),  # For scale, shift, gate for norm1 and norm2
+        )
+        self.txt_norm1 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        # Text doesn't need separate attention - it's handled by img_attn joint computation
+        self.txt_norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=eps)
+        self.txt_mlp = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate")
+
+    def _modulate(self, x, mod_params):
+        """Apply modulation to input tensor"""
+        shift, scale, gate = mod_params.chunk(3, dim=-1)
+        return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1), gate.unsqueeze(1)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        encoder_hidden_states_mask: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Get modulation parameters for both streams
+        img_mod_params = self.img_mod(temb)  # [B, 6*dim]
+        txt_mod_params = self.txt_mod(temb)  # [B, 6*dim]
+
+        # Split modulation parameters for norm1 and norm2
+        img_mod1, img_mod2 = img_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
+        txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1)  # Each [B, 3*dim]
+
+        # Process image stream - norm1 + modulation
+        img_normed = self.img_norm1(hidden_states)
+        img_modulated, img_gate1 = self._modulate(img_normed, img_mod1)
+
+        # Process text stream - norm1 + modulation
+        txt_normed = self.txt_norm1(encoder_hidden_states)
+        txt_modulated, txt_gate1 = self._modulate(txt_normed, txt_mod1)
+
+        # Use QwenAttnProcessor2_0 for joint attention computation
+        # This directly implements the DoubleStreamLayerMegatron logic:
+        # 1. Computes QKV for both streams
+        # 2. Applies QK normalization and RoPE
+        # 3. Concatenates and runs joint attention
+        # 4. Splits results back to separate streams
+        joint_attention_kwargs = joint_attention_kwargs or {}
+        attn_output = self.attn(
+            hidden_states=img_modulated,  # Image stream (will be processed as "sample")
+            encoder_hidden_states=txt_modulated,  # Text stream (will be processed as "context")
+            encoder_hidden_states_mask=encoder_hidden_states_mask,
+            image_rotary_emb=image_rotary_emb,
+            **joint_attention_kwargs,
+        )
+
+        # QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided
+        img_attn_output, txt_attn_output = attn_output
+
+        # Apply attention gates and add residual (like in Megatron)
+        hidden_states = hidden_states + img_gate1 * img_attn_output
+        encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output
+
+        # Process image stream - norm2 + MLP
+        img_normed2 = self.img_norm2(hidden_states)
+        img_modulated2, img_gate2 = self._modulate(img_normed2, img_mod2)
+        img_mlp_output = self.img_mlp(img_modulated2)
+        hidden_states = hidden_states + img_gate2 * img_mlp_output
+
+        # Process text stream - norm2 + MLP
+        txt_normed2 = self.txt_norm2(encoder_hidden_states)
+        txt_modulated2, txt_gate2 = self._modulate(txt_normed2, txt_mod2)
+        txt_mlp_output = self.txt_mlp(txt_modulated2)
+        encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output
+
+        # Clip to prevent overflow for fp16
+        if encoder_hidden_states.dtype == torch.float16:
+            encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+        if hidden_states.dtype == torch.float16:
+            hidden_states = hidden_states.clip(-65504, 65504)
+
+        return encoder_hidden_states, hidden_states
+
+
+class QwenImageTransformer2DModel(ModelMixin, ConfigMixin, PeftAdapterMixin, FromOriginalModelMixin, CacheMixin, AttentionMixin):
+    """
+    The Transformer model introduced in Qwen.
+
+    Args:
+        patch_size (`int`, defaults to `2`):
+            Patch size to turn the input data into small patches.
+        in_channels (`int`, defaults to `64`):
+            The number of channels in the input.
+        out_channels (`int`, *optional*, defaults to `None`):
+            The number of channels in the output. If not specified, it defaults to `in_channels`.
+        num_layers (`int`, defaults to `60`):
+            The number of layers of dual stream DiT blocks to use.
+        attention_head_dim (`int`, defaults to `128`):
+            The number of dimensions to use for each attention head.
+        num_attention_heads (`int`, defaults to `24`):
+            The number of attention heads to use.
+        joint_attention_dim (`int`, defaults to `3584`):
+            The number of dimensions to use for the joint attention (embedding/channel dimension of
+            `encoder_hidden_states`).
+        guidance_embeds (`bool`, defaults to `False`):
+            Whether to use guidance embeddings for guidance-distilled variant of the model.
+        axes_dims_rope (`Tuple[int]`, defaults to `(16, 56, 56)`):
+            The dimensions to use for the rotary positional embeddings.
+    """
+
+    _supports_gradient_checkpointing = True
+    _no_split_modules = ["QwenImageTransformerBlock"]
+    _skip_layerwise_casting_patterns = ["pos_embed", "norm"]
+    _repeated_blocks = ["QwenImageTransformerBlock"]
+
+    @register_to_config
+    def __init__(
+        self,
+        patch_size: int = 2,
+        in_channels: int = 64,
+        out_channels: Optional[int] = 16,
+        num_layers: int = 60,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 24,
+        joint_attention_dim: int = 3584,
+        guidance_embeds: bool = False,  # TODO: this should probably be removed
+        axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
+    ):
+        super().__init__()
+        self.out_channels = out_channels or in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        self.pos_embed = QwenEmbedRope(theta=10000, axes_dim=list(axes_dims_rope), scale_rope=True)
+
+        self.time_text_embed = QwenTimestepProjEmbeddings(embedding_dim=self.inner_dim)
+
+        self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6)
+
+        self.img_in = nn.Linear(in_channels, self.inner_dim)
+        self.txt_in = nn.Linear(joint_attention_dim, self.inner_dim)
+
+        self.transformer_blocks = nn.ModuleList(
+            [
+                QwenImageTransformerBlock(
+                    dim=self.inner_dim,
+                    num_attention_heads=num_attention_heads,
+                    attention_head_dim=attention_head_dim,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6)
+        self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True)
+
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor = None,
+        encoder_hidden_states_mask: torch.Tensor = None,
+        timestep: torch.LongTensor = None,
+        image_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        guidance: torch.Tensor = None,  # TODO: this should probably be removed
+        attention_kwargs: Optional[Dict[str, Any]] = None,
+        return_dict: bool = True,
+    ) -> Union[torch.Tensor, Transformer2DModelOutput]:
+        """
+        The [`QwenTransformer2DModel`] forward method.
+
+        Args:
+            hidden_states (`torch.Tensor` of shape `(batch_size, image_sequence_length, in_channels)`):
+                Input `hidden_states`.
+            encoder_hidden_states (`torch.Tensor` of shape `(batch_size, text_sequence_length, joint_attention_dim)`):
+                Conditional embeddings (embeddings computed from the input conditions such as prompts) to use.
+            encoder_hidden_states_mask (`torch.Tensor` of shape `(batch_size, text_sequence_length)`):
+                Mask of the input conditions.
+            timestep ( `torch.LongTensor`):
+                Used to indicate denoising step.
+            attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~models.transformer_2d.Transformer2DModelOutput`] instead of a plain
+                tuple.
+
+        Returns:
+            If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
+            `tuple` where the first element is the sample tensor.
+        """
+        if attention_kwargs is not None:
+            attention_kwargs = attention_kwargs.copy()
+            lora_scale = attention_kwargs.pop("scale", 1.0)
+        else:
+            lora_scale = 1.0
+
+        if USE_PEFT_BACKEND:
+            # weight the lora layers by setting `lora_scale` for each PEFT layer
+            scale_lora_layers(self, lora_scale)
+        else:
+            if attention_kwargs is not None and attention_kwargs.get("scale", None) is not None:
+                logger.warning(
+                    "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+                )
+
+        hidden_states = self.img_in(hidden_states)
+
+        timestep = timestep.to(hidden_states.dtype)
+        encoder_hidden_states = self.txt_norm(encoder_hidden_states)
+        encoder_hidden_states = self.txt_in(encoder_hidden_states)
+
+        if guidance is not None:
+            guidance = guidance.to(hidden_states.dtype) * 1000
+
+        temb = (
+            self.time_text_embed(timestep, hidden_states)
+            if guidance is None
+            else self.time_text_embed(timestep, guidance, hidden_states)
+        )
+
+        for index_block, block in enumerate(self.transformer_blocks):
+            if torch.is_grad_enabled() and self.gradient_checkpointing:
+                encoder_hidden_states, hidden_states = self._gradient_checkpointing_func(
+                    block,
+                    hidden_states,
+                    encoder_hidden_states,
+                    encoder_hidden_states_mask,
+                    temb,
+                    image_rotary_emb,
+                )
+
+            else:
+                encoder_hidden_states, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_hidden_states_mask=encoder_hidden_states_mask,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    joint_attention_kwargs=attention_kwargs,
+                )
+
+        # Use only the image part (hidden_states) from the dual-stream blocks
+        hidden_states = self.norm_out(hidden_states, temb)
+        output = self.proj_out(hidden_states)
+
+        if USE_PEFT_BACKEND:
+            # remove `lora_scale` from each PEFT layer
+            unscale_lora_layers(self, lora_scale)
+
+        if not return_dict:
+            return (output,)
+
+        return Transformer2DModelOutput(sample=output)

requirements.txt

@@ -1,13 +1,34 @@
-diffusers==0.38.0
-
-transformers
-accelerate
-safetensors
-sentencepiece
-dashscope
-kernels==0.11.0
-torchvision
-peft
-torchao==0.11.0
-spandrel
-spandrel_extra_arches
+# Pro Realism Edit Studio - Enhanced Edition
+# ====================================================
+# Core dependencies
+diffusers==0.38.0
+
+transformers
+accelerate
+safetensors
+sentencepiece
+dashscope
+kernels==0.11.0
+torchvision
+peft
+torchao==0.11.0
+
+# Image processing and upscaling
+spandrel
+spandrel_extra_arches
+Pillow
+numpy
+
+# Face restoration (REQUIRED for enhanced version)
+gfpgan
+
+# OpenCV for face detection (REQUIRED for enhanced version)
+opencv-python
+
+# SciPy for advanced image processing (optional fallback)
+scipy
+
+# Gradio and utilities
+gradio
+spaces
+huggingface_hub

tests/test_enhance_stage.py

@@ -1,45 +1,45 @@
-from pathlib import Path
-import unittest
-
-
-ROOT = Path(__file__).resolve().parents[1]
-
-
-def read_text(relative_path: str) -> str:
-    return (ROOT / relative_path).read_text(encoding="utf-8")
-
-
-class EnhanceStageContractTest(unittest.TestCase):
-    def test_enhance_modes_match_product_contract(self):
-        app = read_text("app.py")
-
-        self.assertIn('ENHANCE_MODE_OFF = "Off"', app)
-        self.assertIn('ENHANCE_MODE_UPSCALE = "Upscale"', app)
-        self.assertIn('ENHANCE_MODE_CLEAN = "Clean"', app)
-        self.assertIn('ENHANCE_MODE_MAX_DETAIL = "Max Detail"', app)
-        self.assertIn(
-            "ENHANCE_MODE_CHOICES = [ENHANCE_MODE_OFF, ENHANCE_MODE_UPSCALE, ENHANCE_MODE_CLEAN, ENHANCE_MODE_MAX_DETAIL]",
-            app,
-        )
-        self.assertIn("enhance_mode = gr.Radio", app)
-        self.assertIn("apply_enhancement", app)
-
-    def test_upscaler_is_lazy_and_configurable(self):
-        app = read_text("app.py")
-
-        self.assertIn("load_upscaler_model", app)
-        self.assertIn("UPSCALER_MODEL_ID", app)
-        self.assertIn("UPSCALER_MODEL_FILENAME", app)
-        self.assertIn("hf_hub_download(", app)
-        self.assertIn("spandrel", app)
-        self.assertIn("tile_upscale", app)
-
-    def test_upscaler_dependencies_are_declared(self):
-        requirements = read_text("requirements.txt")
-
-        self.assertIn("spandrel", requirements)
-        self.assertIn("spandrel_extra_arches", requirements)
-
-
-if __name__ == "__main__":
-    unittest.main()
+from pathlib import Path
+import unittest
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def read_text(relative_path: str) -> str:
+    return (ROOT / relative_path).read_text(encoding="utf-8")
+
+
+class EnhanceStageContractTest(unittest.TestCase):
+    def test_enhance_modes_match_product_contract(self):
+        app = read_text("app.py")
+
+        self.assertIn('ENHANCE_MODE_OFF = "Off"', app)
+        self.assertIn('ENHANCE_MODE_UPSCALE = "Upscale"', app)
+        self.assertIn('ENHANCE_MODE_CLEAN = "Clean"', app)
+        self.assertIn('ENHANCE_MODE_MAX_DETAIL = "Max Detail"', app)
+        self.assertIn(
+            "ENHANCE_MODE_CHOICES = [ENHANCE_MODE_OFF, ENHANCE_MODE_UPSCALE, ENHANCE_MODE_CLEAN, ENHANCE_MODE_MAX_DETAIL]",
+            app,
+        )
+        self.assertIn("enhance_mode = gr.Radio", app)
+        self.assertIn("apply_enhancement", app)
+
+    def test_upscaler_is_lazy_and_configurable(self):
+        app = read_text("app.py")
+
+        self.assertIn("load_upscaler_model", app)
+        self.assertIn("UPSCALER_MODEL_ID", app)
+        self.assertIn("UPSCALER_MODEL_FILENAME", app)
+        self.assertIn("hf_hub_download(", app)
+        self.assertIn("spandrel", app)
+        self.assertIn("tile_upscale", app)
+
+    def test_upscaler_dependencies_are_declared(self):
+        requirements = read_text("requirements.txt")
+
+        self.assertIn("spandrel", requirements)
+        self.assertIn("spandrel_extra_arches", requirements)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_external_sources.py

@@ -1,51 +1,51 @@
-from pathlib import Path
-import unittest
-
-
-ROOT = Path(__file__).resolve().parents[1]
-
-
-def read_text(relative_path: str) -> str:
-    return (ROOT / relative_path).read_text(encoding="utf-8")
-
-
-class ExternalSourcePolicyTest(unittest.TestCase):
-    def test_app_uses_phr00t_v23_single_file_source(self):
-        app = read_text("app.py")
-
-        self.assertNotIn("Sneak-Moose", app)
-        self.assertIn("Phr00t/Qwen-Image-Edit-Rapid-AIO", app)
-        self.assertIn("v23/Qwen-Rapid-AIO-NSFW-v23.safetensors", app)
-        self.assertIn("load_phr00t_rapid_transformer", app)
-        self.assertIn('PHR00T_TRANSFORMER_PREFIX = "model.diffusion_model."', app)
-        self.assertIn("safe_open", app)
-        self.assertNotIn("from_single_file", app)
-
-    def test_private_diagnostics_callback_is_removed(self):
-        app = read_text("app.py")
-
-        self.assertNotIn("QUALITY_ENHANCEMENT_URL", app)
-        self.assertNotIn("QUALITY_ENHANCEMENT_TOKEN", app)
-        self.assertNotIn("requests.post", app)
-        self.assertNotIn("_emit_diagnostics", app)
-
-    def test_third_party_video_space_is_not_hardcoded(self):
-        app = read_text("app.py")
-
-        self.assertNotIn("multimodalart/wan-2-2-first-last-frame", app)
-        self.assertIn('VIDEO_SPACE_ID = os.environ.get("VIDEO_SPACE_ID", "").strip()', app)
-        self.assertIn("Client(VIDEO_SPACE_ID)", app)
-
-    def test_no_external_logo_image_or_sneak_moose_docs(self):
-        relevant_text = "\n".join(
-            read_text(path)
-            for path in ("app.py", "README.md", "AGENT.md", "CLAUDE.md")
-        )
-
-        self.assertNotIn("Sneak-Moose", relevant_text)
-        self.assertNotIn("qianwen-res.oss-cn-beijing.aliyuncs.com", relevant_text)
-        self.assertNotIn("Rapid-AIO v18", relevant_text)
-
-
-if __name__ == "__main__":
-    unittest.main()
+from pathlib import Path
+import unittest
+
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def read_text(relative_path: str) -> str:
+    return (ROOT / relative_path).read_text(encoding="utf-8")
+
+
+class ExternalSourcePolicyTest(unittest.TestCase):
+    def test_app_uses_phr00t_v23_single_file_source(self):
+        app = read_text("app.py")
+
+        self.assertNotIn("Sneak-Moose", app)
+        self.assertIn("Phr00t/Qwen-Image-Edit-Rapid-AIO", app)
+        self.assertIn("v23/Qwen-Rapid-AIO-NSFW-v23.safetensors", app)
+        self.assertIn("load_phr00t_rapid_transformer", app)
+        self.assertIn('PHR00T_TRANSFORMER_PREFIX = "model.diffusion_model."', app)
+        self.assertIn("safe_open", app)
+        self.assertNotIn("from_single_file", app)
+
+    def test_private_diagnostics_callback_is_removed(self):
+        app = read_text("app.py")
+
+        self.assertNotIn("QUALITY_ENHANCEMENT_URL", app)
+        self.assertNotIn("QUALITY_ENHANCEMENT_TOKEN", app)
+        self.assertNotIn("requests.post", app)
+        self.assertNotIn("_emit_diagnostics", app)
+
+    def test_third_party_video_space_is_not_hardcoded(self):
+        app = read_text("app.py")
+
+        self.assertNotIn("multimodalart/wan-2-2-first-last-frame", app)
+        self.assertIn('VIDEO_SPACE_ID = os.environ.get("VIDEO_SPACE_ID", "").strip()', app)
+        self.assertIn("Client(VIDEO_SPACE_ID)", app)
+
+    def test_no_external_logo_image_or_sneak_moose_docs(self):
+        relevant_text = "\n".join(
+            read_text(path)
+            for path in ("app.py", "README.md", "AGENT.md", "CLAUDE.md")
+        )
+
+        self.assertNotIn("Sneak-Moose", relevant_text)
+        self.assertNotIn("qianwen-res.oss-cn-beijing.aliyuncs.com", relevant_text)
+        self.assertNotIn("Rapid-AIO v18", relevant_text)
+
+
+if __name__ == "__main__":
+    unittest.main()