import os
import gc
import gradio as gr
import numpy as np
import spaces
import torch
import random
import base64
import json
import html as html_lib
from io import BytesIO
from PIL import Image
MAX_SEED = np.iinfo(np.int32).max
LANCZOS = getattr(Image, "Resampling", Image).LANCZOS
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
print("torch.__version__ =", torch.__version__)
print("Using device:", device)
from diffusers import FlowMatchEulerDiscreteScheduler
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
dtype = torch.bfloat16
pipe = QwenImageEditPlusPipeline.from_pretrained(
"FireRedTeam/FireRed-Image-Edit-1.1",
transformer=QwenImageTransformer2DModel.from_pretrained(
"prithivMLmods/Qwen-Image-Edit-Rapid-AIO-V19",
torch_dtype=dtype,
device_map="cuda",
),
torch_dtype=dtype,
).to(device)
try:
pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
print("Flash Attention 3 Processor set successfully.")
except Exception as e:
print(f"Warning: Could not set FA3 processor: {e}")
EXAMPLES_CONFIG = [
{
"images": ["examples/1.jpg"],
"prompt": "cinematic polaroid with soft grain subtle vignette gentle lighting white frame handwritten photographed 'Fire-Edit' preserving realistic texture and details.",
},
{
"images": ["examples/2.jpg"],
"prompt": "Transform the image into a dotted cartoon style.",
},
{
"images": ["examples/3.jpeg"],
"prompt": "Convert it to black and white.",
},
{
"images": ["examples/4.jpg", "examples/5.jpg"],
"prompt": "Replace her glasses with the new glasses from image 1.",
},
{
"images": ["examples/8.jpg", "examples/9.png"],
"prompt": "Replace the current clothing with the clothing from the reference image 2. Keep the person's face, hairstyle, body pose, background, lighting, and camera angle unchanged. Ensure the new outfit fits naturally with realistic fabric texture, proper shadows, folds, and accurate proportions. Match the lighting, color tone, and overall style for a seamless and high-quality result.",
},
{
"images": ["examples/10.jpg", "examples/11.png"],
"prompt": "Replace the current clothing with the clothing from the reference image 2. Keep the person's face, hairstyle, body pose, background, lighting, and camera angle unchanged. Ensure the new outfit fits naturally with realistic fabric texture, proper shadows, folds, and accurate proportions. Match the lighting, color tone, and overall style for a seamless and high-quality result.",
},
]
def make_thumb_b64(path, max_dim=220):
if not os.path.exists(path):
return ""
try:
img = Image.open(path).convert("RGB")
img.thumbnail((max_dim, max_dim), LANCZOS)
buf = BytesIO()
img.save(buf, format="JPEG", quality=65)
return f"data:image/jpeg;base64,{base64.b64encode(buf.getvalue()).decode()}"
except Exception as e:
print(f"Thumbnail error for {path}: {e}")
return ""
def encode_full_image(path):
if not os.path.exists(path):
return ""
try:
with open(path, "rb") as f:
data = f.read()
ext = path.rsplit(".", 1)[-1].lower()
mime = {"jpg": "image/jpeg", "jpeg": "image/jpeg", "png": "image/png", "webp": "image/webp"}.get(ext, "image/jpeg")
return f"data:{mime};base64,{base64.b64encode(data).decode()}"
except Exception as e:
print(f"Encode error for {path}: {e}")
return ""
def build_example_cards_html():
cards = ""
for i, ex in enumerate(EXAMPLES_CONFIG):
thumbs_html = ""
for path in ex["images"]:
thumb = make_thumb_b64(path)
if thumb:
thumbs_html += f''
else:
thumbs_html += '
Preview
'
n = len(ex["images"])
badge = f'{n} image{"s" if n > 1 else ""}'
prompt_short = html_lib.escape(ex["prompt"][:90])
if len(ex["prompt"]) > 90:
prompt_short += "..."
cards += f'''