extras/ImageEditServer.py

import argparse
import base64
import io
import time
import torch
import uvicorn
import gc
import asyncio
import traceback
from typing import List, Optional, Union
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, UploadFile, File, Form
from pydantic import BaseModel
from PIL import Image, ImageOps

# Argument parsing
parser = argparse.ArgumentParser(description="Flux Image Edit Server with Nunchaku")
parser.add_argument("--host", type=str, default="0.0.0.0", help="Host to bind to")
parser.add_argument("--port", type=int, default=8000, help="Port to bind to")
parser.add_argument("--model", type=str, default="black-forest-labs/FLUX.1-Kontext-dev", help="Path or Repo ID of the base model")
parser.add_argument("--optimized-model", type=str, default=None, help="Path to the optimized Nunchaku model safetensors file")
parser.add_argument("--optimized-edit-model", type=str, default=None, help="Path to the optimized Nunchaku model safetensors file for editing (optional)")
parser.add_argument("--backend", type=str, default="kontext", choices=["kontext", "flux2", "flux2-klein", "flux2_klein", "qwen", "qwen21", "qwen-2.1", "qwen21-nvfp4", "qwen21_nvfp4", "glm", "zimage"], help="Backend to use: 'kontext', 'flux2', 'flux2-klein', 'qwen', 'qwen21', 'qwen21-nvfp4', 'glm', or 'zimage'")
parser.add_argument("--steps", type=int, default=28, help="Default number of inference steps")
parser.add_argument("--guidance-scale", type=float, default=3.5, help="Default guidance scale")
parser.add_argument("--qwenimage", action="store_true", help="Use QwenImageBackend (T2I only) instead of full Qwen edit backend")
parser.add_argument("--uma", action="store_true", help="Enable Unified Memory Architecture mode (load all to GPU, disable offload)")
parser.add_argument("--no-layerwise-offload", action="store_true", help="Keep transformer resident in VRAM without layer-by-layer offloading")
parser.add_argument(
    "--text-encoder",
    type=str,
    default=None,
    help="Path or Repo ID of custom text encoder (e.g. models/Qwen_Text_Encoder_Heretic)",
)
parser.add_argument(
    "--nvfp4-text-encoder",
    type=str,
    default=None,
    help=(
        "Path to an NVFP4-pack-quantized HuggingFace text encoder "
        "swaps in vLLM's W4A4 NVFP4 CUTLASS GEMM for ~4x text-encoder VRAM savings."
    ),
)
args = parser.parse_args()

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup logic
    load_model()
    yield
    # Shutdown logic (if any) could go here

app = FastAPI(lifespan=lifespan)

# Global components
IMAGE_DIMENSION_ALIGNMENT = 32
# Cap for a reference image fed with input_fit=native: the pipeline resizes each condition image
# at its own aspect ratio, so this only bounds token count, never framing.
MAX_REFERENCE_PIXELS = 1024 * 1024
pipeline = None
edit_pipeline = None
request_lock = asyncio.Lock()
is_sleeping_flag = False
sleep_requested = False

def set_zero_cond_t(transformer, val: bool):
    """Recursively synchronize zero_cond_t across transformer and all inner blocks."""
    if transformer is None:
        return
    if hasattr(transformer, "zero_cond_t"):
        transformer.zero_cond_t = val
    if hasattr(transformer, "transformer_blocks"):
        for block in transformer.transformer_blocks:
            if hasattr(block, "zero_cond_t"):
                block.zero_cond_t = val

def load_model():
    global pipeline, edit_pipeline
    
    try:
        if args.backend == "kontext":
            import KontextBackend
            print(f"Initializing KontextBackend...")
            backend = KontextBackend.KontextBackend(args.model, args.optimized_model)
            pipeline, edit_pipeline = backend.load()
        elif args.backend == "flux2":
            import Flux2Backend
            print(f"Initializing Flux2Backend...")
            backend = Flux2Backend.Flux2Backend(args.model)
            pipeline, edit_pipeline = backend.load()
        elif args.backend == "glm":
            import GlmBackend
            print(f"Initializing GlmBackend...")
            # Use provided model or default to the one in the snippet if args.model is generic
            # The user might pass the specific GLM model via --model, or we default in GlmBackend.
            # Let's pass args.model if it's not the default flux one, otherwise let GlmBackend use its default.
            model_to_use = args.model if args.model != "black-forest-labs/FLUX.1-Kontext-dev" else "Disty0/GLM-Image-SDNQ-4bit-dynamic"
            backend = GlmBackend.GlmBackend(model_to_use)
            pipeline, edit_pipeline = backend.load()
        elif args.backend in ["qwen21", "qwen-2.1", "qwen21-nvfp4", "qwen21_nvfp4"]:
            if args.optimized_model or "nvfp4" in args.backend:
                import QwenImage21NVFP4Backend
                print(f"Initializing QwenImage21NVFP4Backend (resident NVFP4)...")
                backend = QwenImage21NVFP4Backend.QwenImage21NVFP4Backend(
                    model_id=args.model,
                    optimized_model_path=args.optimized_model or "/home/olegk/Nikola/models/nunchaku-qwen-image-2.1/best_quality_fp4.safetensors",
                    text_encoder_path=args.text_encoder,
                )
                pipeline, edit_pipeline = backend.load()
            else:
                import QwenImage21Backend
                print(f"Initializing QwenImage21Backend (unquantized baseline)...")
                backend = QwenImage21Backend.QwenImage21Backend(
                    args.model,
                    text_encoder_path=args.text_encoder,
                )
                pipeline, edit_pipeline = backend.load()
        elif args.backend.startswith("qwen"):
            if args.qwenimage:
                import QwenImageBackend
                print(f"Initializing QwenImageBackend (T2I only)...")
                backend = QwenImageBackend.QwenImageBackend(args.model, args.optimized_model)
                pipeline, edit_pipeline = backend.load()
            else:
                import QwenBackend
                print(f"Initializing QwenBackend...")
                backend = QwenBackend.QwenBackend(
                    args.model,
                    args.optimized_model,
                    optimized_edit_model_path=args.optimized_edit_model,
                    uma=args.uma,
                    no_layerwise_offload=args.no_layerwise_offload,
                    text_encoder_path=args.text_encoder,
                )
                pipeline, edit_pipeline = backend.load()
        elif args.backend in ["flux2-klein", "flux2_klein"]:
            import Flux2KleinBackend
            print(f"Initializing Flux2KleinBackend...")
            backend = Flux2KleinBackend.Flux2KleinBackend(
                args.model,
                args.optimized_model,
                nvfp4_text_encoder_path=args.nvfp4_text_encoder,
            )
            pipeline, edit_pipeline = backend.load()
        elif args.backend == "zimage":
            import ZImageTurboBackend
            print(f"Initializing ZImageTurboBackend...")
            backend = ZImageTurboBackend.ZImageTurboBackend(
                args.model,
                args.optimized_model,
                uma=args.uma,
                nvfp4_text_encoder_path=args.nvfp4_text_encoder,
            )
            pipeline, edit_pipeline = backend.load()
        else:
            raise ValueError(f"Unknown backend: {args.backend}")
            
    except Exception as e:
        print(f"Oh no! The model refused to wake up: {e}")
        raise e
    
    # Enable progress bar for diffusers
    import diffusers.utils.logging
    diffusers.utils.logging.enable_progress_bar()
    diffusers.utils.logging.set_verbosity_info()
    
    # Instrument timing tracker for stage-by-stage profiling
    instrument_pipeline(pipeline)
    instrument_pipeline(edit_pipeline)
    
    print("Model loaded successfully! Ready for editing quests!")

class PipelineTimingTracker:
    def __init__(self):
        self.reset()

    def reset(self):
        self.text_encoder_time = 0.0
        self.transformer_times = []
        self.vae_time = 0.0
        self._te_start = None
        self._tr_start = None
        self._vae_start = None

    def on_te_start(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        self._te_start = time.perf_counter()

    def on_te_end(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        if self._te_start is not None:
            self.text_encoder_time += time.perf_counter() - self._te_start
            self._te_start = None

    def on_tr_start(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        self._tr_start = time.perf_counter()

    def on_tr_end(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        if self._tr_start is not None:
            self.transformer_times.append(time.perf_counter() - self._tr_start)
            self._tr_start = None

    def on_vae_start(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        self._vae_start = time.perf_counter()

    def on_vae_end(self):
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        if self._vae_start is not None:
            self.vae_time += time.perf_counter() - self._vae_start
            self._vae_start = None

    def get_summary(self, total_wall_time: float) -> dict:
        num_steps = len(self.transformer_times)
        total_tr = sum(self.transformer_times)
        avg_step = (total_tr / num_steps) if num_steps > 0 else 0.0
        return {
            "text_encoder_s": round(self.text_encoder_time, 3),
            "transformer_s": round(total_tr, 3),
            "steps": num_steps,
            "step_avg_s": round(avg_step, 3),
            "vae_s": round(self.vae_time, 3),
            "total_s": round(total_wall_time, 3),
        }

timing_tracker = PipelineTimingTracker()

def instrument_pipeline(p):
    if p is None or getattr(p, "_timing_instrumented", False):
        return
    if hasattr(p, "text_encoder") and p.text_encoder is not None:
        orig_te = p.text_encoder.forward
        def timed_te(*args, **kwargs):
            timing_tracker.on_te_start()
            try:
                return orig_te(*args, **kwargs)
            finally:
                timing_tracker.on_te_end()
        p.text_encoder.forward = timed_te

    if hasattr(p, "transformer") and p.transformer is not None:
        orig_tr = p.transformer.forward
        def timed_tr(*args, **kwargs):
            timing_tracker.on_tr_start()
            try:
                return orig_tr(*args, **kwargs)
            finally:
                timing_tracker.on_tr_end()
        p.transformer.forward = timed_tr

    if hasattr(p, "vae") and p.vae is not None:
        orig_vae = p.vae.decode
        def timed_vae(*args, **kwargs):
            timing_tracker.on_vae_start()
            try:
                return orig_vae(*args, **kwargs)
            finally:
                timing_tracker.on_vae_end()
        p.vae.decode = timed_vae

    p._timing_instrumented = True

def flush():
    gc.collect()
    torch.cuda.empty_cache()


class ImageGenerationRequest(BaseModel):
    prompt: str
    n: int = 1
    size: str = "1024x1024"
    response_format: str = "b64_json"
    quality: str = "standard"
    style: str = "vivid"
    num_inference_steps: Optional[int] = None
    guidance_scale: Optional[float] = None
    negative_prompt: Optional[str] = None
    seed: Optional[int] = None


@app.post("/v1/sleep")
async def sleep_endpoint():
    global is_sleeping_flag, sleep_requested
    sleep_requested = True
    try:
        async with request_lock:
            if not is_sleeping_flag and sleep_requested:
                print("Sleep requested, moving models to CPU...")
                for p in [pipeline, edit_pipeline]:
                    if not p: continue
                    for name, component in p.components.items():
                        if isinstance(component, torch.nn.Module):
                            # Special handling for Nunchaku which blocks .to() if offload is True
                            if hasattr(component, "set_offload") and getattr(component, "offload", False):
                                component.set_offload(False)
                                component._nunchaku_was_offloaded = True

                            try:
                                component.to("cpu")
                            except Exception as e:
                                pass
                flush()
                is_sleeping_flag = True
    finally:
        sleep_requested = False
    return {"status": "sleep completed", "is_sleeping": is_sleeping_flag}

@app.post("/v1/wake_up")
async def wake_up_endpoint():
    global is_sleeping_flag, sleep_requested
    sleep_requested = False
    async with request_lock:
        if is_sleeping_flag:
            print("Waking up, restoring models to CUDA...")
            for p in [pipeline, edit_pipeline]:
                if not p: continue
                excluded = getattr(p, "_exclude_from_cpu_offload", [])
                for name, component in p.components.items():
                    if isinstance(component, torch.nn.Module):
                        if getattr(component, "_nunchaku_was_offloaded", False):
                            component.set_offload(True, use_pin_memory=True, num_blocks_on_gpu=8)
                            for attr in ["img_in", "txt_in", "txt_norm", "time_text_embed", "norm_out", "proj_out"]:
                                if hasattr(component, attr):
                                    try:
                                        getattr(component, attr).to("cuda")
                                    except Exception:
                                        pass
                            component._nunchaku_was_offloaded = False
                        elif not hasattr(component, "_hf_hook") or name in excluded:
                            try:
                                component.to("cuda")
                            except Exception:
                                pass
            is_sleeping_flag = False
    return {"status": "awoken", "is_sleeping": False}

@app.get("/v1/is_sleeping")
async def is_sleeping_endpoint():
    return {"is_sleeping": is_sleeping_flag}


@app.get("/v1/memory_stats")
async def memory_stats_endpoint():
    """Lightweight introspection endpoint that returns PyTorch's CUDA allocator
    snapshot. Used to diagnose VRAM/UMA bloat without restarting the server."""
    stats = {}
    if torch.cuda.is_available():
        stats["allocated_gb"] = torch.cuda.memory_allocated() / 1e9
        stats["reserved_gb"] = torch.cuda.memory_reserved() / 1e9
        stats["max_allocated_gb"] = torch.cuda.max_memory_allocated() / 1e9
        stats["max_reserved_gb"] = torch.cuda.max_memory_reserved() / 1e9
        # Top allocations by size from the allocator snapshot (>=64 MiB)
        try:
            snap = torch.cuda.memory_snapshot()
            blocks = []
            for seg in snap:
                for b in seg.get("blocks", []):
                    if b.get("state") == "active_allocated" and b.get("size", 0) >= 64 * 1024 * 1024:
                        blocks.append(b["size"])
            blocks.sort(reverse=True)
            stats["large_active_blocks_gb"] = [round(s / 1e9, 3) for s in blocks[:20]]
            stats["large_active_blocks_total_gb"] = round(sum(blocks) / 1e9, 3)
            stats["large_active_blocks_count"] = len(blocks)
        except Exception as e:
            stats["snapshot_error"] = str(e)
        # Walk Python objects to find big tensors and group them
        try:
            import gc as _gc
            seen = set()
            big = []
            for obj in _gc.get_objects():
                try:
                    if isinstance(obj, torch.Tensor) and obj.is_cuda:
                        ptr = obj.data_ptr()
                        if ptr in seen or ptr == 0:
                            continue
                        seen.add(ptr)
                        sz = obj.element_size() * obj.numel()
                        if sz >= 16 * 1024 * 1024:
                            big.append((sz, tuple(obj.shape), str(obj.dtype)))
                except Exception:
                    continue
            big.sort(reverse=True)
            # Group by (shape, dtype)
            from collections import Counter
            grouped = Counter((shape, dtype) for _, shape, dtype in big)
            stats["big_tensor_groups"] = [
                {"shape": list(shape), "dtype": dtype, "count": cnt,
                 "size_gb_each": round(
                     (1 if shape == () else (lambda l: __import__('functools').reduce(lambda a, b: a*b, l, 1))(shape)) * (
                         8 if 'int64' in dtype or 'float64' in dtype else
                         4 if 'int32' in dtype or 'float32' in dtype else
                         2 if 'bfloat16' in dtype or 'float16' in dtype else 1
                     ) / 1e9, 4)}
                for (shape, dtype), cnt in grouped.most_common(30)
            ]
            stats["big_tensor_count"] = len(big)
            stats["big_tensor_total_gb"] = round(sum(s for s, _, _ in big) / 1e9, 3)
        except Exception as e:
            stats["walk_error"] = str(e)
    return stats

@app.post("/v1/images/edits")
async def edit_image(
    image: Union[List[UploadFile], UploadFile] = File(...),
    prompt: str = Form(...),
    n: int = Form(1),
    size: str = Form("1024x1024"),
    width: Optional[int] = Form(None),
    height: Optional[int] = Form(None),
    input_fit: str = Form("canvas"),
    response_format: str = Form("b64_json"), # Default to b64_json
    guidance_scale: Optional[float] = Form(None), 
    num_inference_steps: Optional[int] = Form(None),
    negative_prompt: Optional[str] = Form(None),
    seed: Optional[int] = Form(None)
):
    # Use CLI defaults if not provided
    steps = num_inference_steps if num_inference_steps is not None else args.steps
    cfg_scale = guidance_scale if guidance_scale is not None else args.guidance_scale
    neg_prompt = negative_prompt if negative_prompt is not None else "" # Default empty for now, or maybe None?
    
    generator = None
    import random
    if seed is None:
        seed = random.randint(0, 2**32 - 1)
    
    print(f"Using seed: {seed}")
    generator = torch.Generator(device="cuda").manual_seed(seed)

    if not edit_pipeline:
        raise HTTPException(status_code=500, detail="Model not loaded")

    if sleep_requested or is_sleeping_flag:
        raise HTTPException(status_code=503, detail="Server is sleeping or trying to sleep.")

    async with request_lock:
        print(f"Received edit request: {prompt}")

        # Processing the input image(s)
        input_files = image if isinstance(image, list) else [image]
        init_images = []

        try:
            for img_file in input_files:
                await img_file.seek(0)
                contents = await img_file.read()
                img = Image.open(io.BytesIO(contents)).convert("RGB")
                init_images.append(img)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid image file: {e}")
        
        if not init_images:
             raise HTTPException(status_code=400, detail="No images provided")

        # Parse max target dimensions from requested size
        try:
            target_width, target_height = map(int, size.split("x"))
        except ValueError:
            target_width, target_height = 1024, 1024

        # Calculate new dimensions preserving aspect ratio based on the first image
        if width is not None and height is not None:
            # Explicit output canvas: the caller owns the framing, so references of any aspect
            # ratio can be fed without dictating the composition's shape.
            req_width, req_height = width, height
        else:
            first_image = init_images[0]
            orig_width, orig_height = first_image.size
            scale = min(target_width / orig_width, target_height / orig_height)
            req_width = int(orig_width * scale)
            req_height = int(orig_height * scale)

        # Ensure dimensions are aligned to 32 for compatibility (e.g. GLM-Image)
        width = (req_width // IMAGE_DIMENSION_ALIGNMENT) * IMAGE_DIMENSION_ALIGNMENT
        height = (req_height // IMAGE_DIMENSION_ALIGNMENT) * IMAGE_DIMENSION_ALIGNMENT

        if (input_fit or "").strip().lower() == "native":
            # QwenImageEditPlus resizes every condition image at its own aspect ratio, so padding
            # references onto the output canvas only costs them resolution and adds false letterbox
            # content the model then reproduces. Keep each reference's own framing.
            resized_images = []
            for img in init_images:
                iw, ih = img.size
                scale = min(1.0, (MAX_REFERENCE_PIXELS / float(max(1, iw * ih))) ** 0.5)
                if scale < 1.0:
                    nw = max(IMAGE_DIMENSION_ALIGNMENT,
                             (int(iw * scale) // IMAGE_DIMENSION_ALIGNMENT) * IMAGE_DIMENSION_ALIGNMENT)
                    nh = max(IMAGE_DIMENSION_ALIGNMENT,
                             (int(ih * scale) // IMAGE_DIMENSION_ALIGNMENT) * IMAGE_DIMENSION_ALIGNMENT)
                    img = img.resize((nw, nh), Image.LANCZOS)
                resized_images.append(img)
        else:
            # Resize input images to match the calculated target size, padding if necessary
            resized_images = []
            for img in init_images:
                if img.size != (width, height):
                    # This handles cases where subsequent images might have different ARs
                    img = ImageOps.pad(img, (width, height), method=Image.LANCZOS, color=(0, 0, 0))
                resized_images.append(img)
            
        # If single image, pass as item, if multiple, pass as list
        # GLM pipeline has a bug where it checks len() on the input, so it must be a list
        if len(resized_images) > 1 or args.backend == "glm":
            image_input = resized_images
        else:
            image_input = resized_images[0]

        response_images = []
        timing_tracker.reset()
        t_req_start = time.perf_counter()
        
        try:
            if args.backend.startswith("qwen"):
                # Qwen specific parameters
                if hasattr(edit_pipeline, "transformer"):
                    set_zero_cond_t(edit_pipeline.transformer, True)
                # guidance_scale maps to true_cfg_scale
                if args.qwenimage: # QwenImageBackend is T2I only, so it doesn't take an image
                    generated_images = edit_pipeline(
                        prompt=prompt,
                        height=height,
                        width=width,
                        num_inference_steps=steps,
                        true_cfg_scale=cfg_scale,
                        num_images_per_prompt=n,
                        generator=generator,
                    ).images
                else: # Full Qwen edit backend takes an image (or list of images now)
                    generated_images = edit_pipeline(
                        image=image_input,
                        prompt=prompt,
                        height=height,
                        width=width,
                        negative_prompt=neg_prompt,
                        num_inference_steps=steps,
                        true_cfg_scale=cfg_scale,
                        num_images_per_prompt=n,
                        generator=generator,
                    ).images
            else:
                # Standard Flux/Kontext or GLM
                # GLM I2I Fix: Manually move vision encoder to GPU because get_image_features escapes hooks
                if args.backend == "glm" and hasattr(edit_pipeline, "vision_language_encoder"):
                     print("Manually moving GLM Vision Encoder to GPU...")
                     edit_pipeline.vision_language_encoder.to("cuda")
                
                try:
                    generated_images = edit_pipeline(
                        image=image_input,
                        prompt=prompt,
                        height=height,
                        width=width,
                        num_inference_steps=steps, 
                        guidance_scale=cfg_scale,
                        num_images_per_prompt=n,
                        generator=generator,
                    ).images
                finally:
                    if args.backend == "glm" and hasattr(edit_pipeline, "vision_language_encoder"):
                         print("Moving GLM Vision Encoder back to CPU...")
                         edit_pipeline.vision_language_encoder.to("cpu")
            
            for img in generated_images:
                buffered = io.BytesIO()
                img.save(buffered, format="PNG")
                img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
                
                if response_format == "b64_json":
                    response_images.append({"b64_json": img_str})
                else:
                    # If url is requested we can't really do it without storage, so we fallback or error?
                    # For now, let's just assume simple b64_json as per request
                    response_images.append({"b64_json": img_str}) # Fallback
                
        except Exception as e:
            print(f"Error during editing: {e}")
            print(traceback.format_exc())
            raise HTTPException(status_code=500, detail=str(e))
        finally:
            flush()

        t_req_total = time.perf_counter() - t_req_start
        timings = timing_tracker.get_summary(t_req_total)
        print(f"[TIMING] Text Encoder: {timings['text_encoder_s']:.2f}s | "
              f"Transformer ({timings['steps']} steps): {timings['transformer_s']:.2f}s ({timings['step_avg_s']:.2f}s/step) | "
              f"VAE: {timings['vae_s']:.2f}s | Total: {timings['total_s']:.2f}s")

        return {
            "created": int(time.time()),
            "timings": timings,
            "seed": seed,
            "width": width,
            "height": height,
            "data": response_images
        }



@app.post("/v1/images/generations")
async def generate_image(request: ImageGenerationRequest):
    if not pipeline:
        raise HTTPException(status_code=500, detail="Model not loaded")

    if sleep_requested or is_sleeping_flag:
        raise HTTPException(status_code=503, detail="Server is sleeping or trying to sleep.")

    async with request_lock:
        #print(f"Received generation request: {request.prompt}")

        # Parse size
        try:
            width, height = map(int, request.size.split("x"))
        except ValueError:
            width, height = 1024, 1024

        # Ensure dimensions are aligned to 32
        width = (width // IMAGE_DIMENSION_ALIGNMENT) * IMAGE_DIMENSION_ALIGNMENT
        height = (height // IMAGE_DIMENSION_ALIGNMENT) * IMAGE_DIMENSION_ALIGNMENT

        response_images = []
        timing_tracker.reset()
        t_req_start = time.perf_counter()
        
        try:
            # Generate images (no image argument for txt2img!)
            steps = request.num_inference_steps if request.num_inference_steps is not None else args.steps
            cfg_scale = request.guidance_scale if request.guidance_scale is not None else args.guidance_scale
            # negative_prompt not in standard request body in original snippet, but we added it to model
            neg_prompt = request.negative_prompt if request.negative_prompt is not None else ""
            
            generator = None
            import random
            seed = request.seed
            if seed is None:
                seed = random.randint(0, 2**32 - 1)
            
            print(f"Using seed: {seed}")
            generator = torch.Generator(device="cuda").manual_seed(seed)

            if args.backend.startswith("qwen"):
                if hasattr(pipeline, "transformer"):
                    set_zero_cond_t(pipeline.transformer, False)
                generated_images = pipeline(
                    prompt=request.prompt,
                    height=height,
                    width=width,
                    num_inference_steps=steps, 
                    true_cfg_scale=cfg_scale,     
                    num_images_per_prompt=request.n,
                    negative_prompt=neg_prompt,
                    generator=generator,
                ).images
            else:
                generated_images = pipeline(
                    prompt=request.prompt,
                    height=height,
                    width=width,
                    num_inference_steps=steps, 
                    guidance_scale=cfg_scale,     
                    num_images_per_prompt=request.n,
                    generator=generator,
                    # Not passing negative_prompt here for generation unless we confirm support in standard Flux pipeline?
                ).images
            
            for img in generated_images:
                buffered = io.BytesIO()
                img.save(buffered, format="PNG")
                img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
                response_images.append({"b64_json": img_str})
                
        except Exception as e:
            print(f"Error during generation: {e}")
            print(traceback.format_exc())
            raise HTTPException(status_code=500, detail=str(e))
        finally:
            flush()

        t_req_total = time.perf_counter() - t_req_start
        timings = timing_tracker.get_summary(t_req_total)
        print(f"[TIMING] Text Encoder: {timings['text_encoder_s']:.2f}s | "
              f"Transformer ({timings['steps']} steps): {timings['transformer_s']:.2f}s ({timings['step_avg_s']:.2f}s/step) | "
              f"VAE: {timings['vae_s']:.2f}s | Total: {timings['total_s']:.2f}s")

        return {
            "created": int(time.time()),
            "timings": timings,
            "seed": seed,
            "width": width,
            "height": height,
            "data": response_images
        }

if __name__ == "__main__":
    uvicorn.run(app, host=args.host, port=args.port)