#!/usr/bin/env python3 """ v1.0: Direct image generation bypassing ComfyUI GGUF OSError Uses Hugging Face diffusers + Qwen Image model directly. """ import sys import os import json import io import base64 import tempfile import time import traceback from pathlib import Path # Configuration MODEL_ID = "Qwen/Qwen-Image" HF_CACHE = str(Path(os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface")))) OUTPUT_DIR = Path(__file__).parent / "outputs" / "images" OUTPUT_DIR.mkdir(parents=True, exist_ok=True) def generate_image(prompt, negative_prompt="", width=1024, height=1024, steps=8, cfg=4.5, seed=None, output_path=None): """ Generate image using Qwen-Image diffusers pipeline. Falls back to CPU if CUDA OOM. """ import random if seed is None: seed = random.randint(0, 2**32 - 1) try: from diffusers import QwenImagePipeline import torch print(f"[ImageGen] Loading Qwen-Image pipeline...", flush=True) # Determine device if torch.cuda.is_available(): device = "cuda" torch_dtype = torch.bfloat16 print(f"[ImageGen] Using CUDA: {torch.cuda.get_device_name(0)} ({torch.cuda.get_device_properties(0).total_mem/1e9:.1f}GB)", flush=True) else: device = "cpu" torch_dtype = torch.float32 print(f"[ImageGen] Using CPU", flush=True) # Load pipeline with optimizations pipe = QwenImagePipeline.from_pretrained( MODEL_ID, torch_dtype=torch_dtype, cache_dir=HF_CACHE, ) # Move to device if device == "cuda": try: pipe = pipe.to(device) if hasattr(pipe, 'enable_model_cpu_offload'): pipe.enable_model_cpu_offload() print(f"[ImageGen] CPU offload enabled", flush=True) elif hasattr(pipe, 'enable_sequential_cpu_offload'): pipe.enable_sequential_cpu_offload() print(f"[ImageGen] Sequential CPU offload enabled", flush=True) except Exception as e: print(f"[ImageGen] GPU load failed, falling back to CPU: {e}", flush=True) pipe = pipe.to("cpu") print(f"[ImageGen] Generating: {prompt[:80]}... ({width}x{height}, {steps} steps)", flush=True) t0 = time.time() # Generate with torch.inference_mode(): image = pipe( prompt=prompt, negative_prompt=negative_prompt or None, width=width, height=height, num_inference_steps=steps, guidance_scale=cfg, generator=torch.Generator(device=device if device == "cuda" else "cpu").manual_seed(seed), ).images[0] elapsed = time.time() - t0 print(f"[ImageGen] Done in {elapsed:.1f}s", flush=True) # Save if output_path is None: ts = int(time.time()) output_path = str(OUTPUT_DIR / f"qwen_gen_{ts}_{seed}.png") image.save(output_path, format="PNG") print(f"[ImageGen] Saved to: {output_path}", flush=True) # Also return as bytes buf = io.BytesIO() image.save(buf, format="PNG") img_bytes = buf.getvalue() return { "success": True, "path": output_path, "image_base64": base64.b64encode(img_bytes).decode("ascii"), "seed": seed, "elapsed": round(elapsed, 1), "size": f"{width}x{height}", } except ImportError as e: print(f"[ImageGen] diffusers not installed: {e}", flush=True) return {"success": False, "error": f"diffusers/QwenImagePipeline not available: {e}"} except Exception as e: traceback.print_exc() print(f"[ImageGen] Failed: {e}", flush=True) return {"success": False, "error": str(e)} def generate_image_fallback(prompt, negative_prompt="", width=1024, height=1024, steps=4, cfg=4.5, seed=None, output_path=None): """ Fallback: Use ComfyUI API (if it's working). """ import urllib.request import urllib.error comfyui_url = "http://127.0.0.1:8188" # Build workflow import random if seed is None: seed = random.randint(0, 2**32 - 1) workflow = { "1": {"class_type": "UnetLoaderGGUF", "inputs": {"unet_name": "qwen-image-Q4_K_M.gguf"}}, "2": {"class_type": "CLIPLoader", "inputs": {"clip_name": "qwen_2.5_vl_7b_fp8_scaled.safetensors", "type": "qwen_image", "device": "default"}}, "3": {"class_type": "VAELoader", "inputs": {"vae_name": "qwen_image_vae.safetensors"}}, "4": {"class_type": "CLIPTextEncode", "inputs": {"text": prompt, "clip": ["2", 0]}}, "5": {"class_type": "CLIPTextEncode", "inputs": {"text": negative_prompt or "", "clip": ["2", 0]}}, "8": {"class_type": "EmptyLatentImage", "inputs": {"width": width, "height": height, "batch_size": 1}}, "9": {"class_type": "KSampler", "inputs": {"seed": seed, "steps": steps, "cfg": cfg, "sampler_name": "euler", "scheduler": "simple", "denoise": 1, "model": ["1", 0], "positive": ["4", 0], "negative": ["5", 0], "latent_image": ["8", 0]}}, "10": {"class_type": "VAEDecode", "inputs": {"samples": ["9", 0], "vae": ["3", 0]}}, "11": {"class_type": "SaveImage", "inputs": {"filename_prefix": "Fallback", "images": ["10", 0]}}, } try: req = urllib.request.Request( f"{comfyui_url}/prompt", data=json.dumps({"prompt": workflow}).encode(), headers={"Content-Type": "application/json"} ) with urllib.request.urlopen(req, timeout=30) as resp: result = json.loads(resp.read()) prompt_id = result.get("prompt_id") if not prompt_id: return {"success": False, "error": "No prompt_id"} # Wait for completion (max 120s) for _ in range(60): time.sleep(2) try: hist_req = urllib.request.Request(f"{comfyui_url}/history/{prompt_id}") with urllib.request.urlopen(hist_req, timeout=10) as resp: hist = json.loads(resp.read()) if prompt_id in hist: entry = hist[prompt_id] status = entry.get("status", {}) if status.get("status_str") == "error": msgs = status.get("messages", []) err = msgs[-1][0] if msgs else "Unknown error" return {"success": False, "error": err} if status.get("completed"): outputs = entry.get("outputs", {}) for node_id, output in outputs.items(): images = output.get("images", []) if images: filename = images[0]["filename"] subfolder = images[0].get("subfolder", "") img_url = f"{comfyui_url}/view?filename={urllib.request.quote(filename)}&subfolder={urllib.request.quote(subfolder)}&type=output" img_req = urllib.request.Request(img_url) with urllib.request.urlopen(img_req, timeout=30) as img_resp: img_bytes = img_resp.read() if output_path: with open(output_path, 'wb') as f: f.write(img_bytes) return { "success": True, "path": str(output_path) if output_path else None, "image_base64": base64.b64encode(img_bytes).decode("ascii"), "seed": seed, "elapsed": _ * 2, "size": f"{width}x{height}", } except Exception: continue return {"success": False, "error": "Timeout waiting for ComfyUI"} except Exception as e: return {"success": False, "error": str(e)} if __name__ == "__main__": # CLI mode: python image_gen.py "prompt" [width] [height] [steps] [cfg] [seed] import argparse parser = argparse.ArgumentParser() parser.add_argument("prompt", help="Image generation prompt") parser.add_argument("--negative", default="", help="Negative prompt") parser.add_argument("--width", type=int, default=1024) parser.add_argument("--height", type=int, default=1024) parser.add_argument("--steps", type=int, default=8) parser.add_argument("--cfg", type=float, default=4.5) parser.add_argument("--seed", type=int, default=None) parser.add_argument("--output", default=None, help="Output path") parser.add_argument("--comfyui", action="store_true", help="Use ComfyUI fallback") args = parser.parse_args() if args.comfyui: result = generate_image_fallback( args.prompt, args.negative, args.width, args.height, args.steps, args.cfg, args.seed, args.output ) else: result = generate_image( args.prompt, args.negative, args.width, args.height, args.steps, args.cfg, args.seed, args.output ) print(json.dumps(result, ensure_ascii=False, indent=2))