Spaces:
Running on Zero
Running on Zero
File size: 2,937 Bytes
f2ec79c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | """Local Qwen Image Edit backend (needs a CUDA GPU).
Mirrors ``linoyts/Qwen-Image-Edit-Angles``: the base Qwen Image Edit 2509
pipeline with Phr00t's Rapid-AIO transformer for fast 4-step inference and
dx8152's multiple-angles LoRA fused in. Heavy imports are deferred so the
rest of the app works on CPU-only machines.
"""
from __future__ import annotations
import random
from typing import Optional
from PIL import Image
from ..config import (
ANGLES_LORA_REPO,
ANGLES_LORA_WEIGHT,
QWEN_BASE_MODEL,
QWEN_RAPID_TRANSFORMER,
)
from ..images import fit_image
from .base import ImageEditBackend
class LocalQwenBackend(ImageEditBackend):
source = "local_qwen_image_edit"
def __init__(self, image_size: int, lora_scale: float = 1.25) -> None:
self.image_size = image_size
self.lora_scale = lora_scale
self._pipe = None
self._device = "cpu"
def prepare(self) -> None:
import torch # local import: heavy dep
from diffusers import QwenImageEditPlusPipeline, QwenImageTransformer2DModel
if not torch.cuda.is_available():
raise RuntimeError("Local Qwen backend requires a CUDA GPU.")
self._device = "cuda"
dtype = torch.bfloat16
transformer = QwenImageTransformer2DModel.from_pretrained(
QWEN_RAPID_TRANSFORMER,
subfolder="transformer",
torch_dtype=dtype,
device_map="cuda",
)
pipe = QwenImageEditPlusPipeline.from_pretrained(
QWEN_BASE_MODEL,
transformer=transformer,
torch_dtype=dtype,
).to(self._device)
pipe.load_lora_weights(
ANGLES_LORA_REPO,
weight_name=ANGLES_LORA_WEIGHT,
adapter_name="angles",
)
pipe.set_adapters(["angles"], adapter_weights=[1.0])
pipe.fuse_lora(adapter_names=["angles"], lora_scale=self.lora_scale)
pipe.unload_lora_weights()
self._pipe = pipe
def edit(
self,
image: Image.Image,
prompt: str,
seed: int,
num_inference_steps: int,
true_guidance_scale: float,
) -> Image.Image:
import torch # local import: heavy dep
base = fit_image(image.convert("RGB"), self.image_size)
if not prompt.strip():
return base
if self._pipe is None:
raise RuntimeError("Backend not prepared; call prepare() first.")
generator = torch.Generator(device=self._device).manual_seed(int(seed))
result = self._pipe(
image=[base],
prompt=prompt,
width=base.width,
height=base.height,
num_inference_steps=num_inference_steps,
generator=generator,
true_cfg_scale=true_guidance_scale,
num_images_per_prompt=1,
).images[0]
return fit_image(result.convert("RGB"), self.image_size)
|