import os

# 防御：某些容器 no_proxy 里有 [::1] 这种括号 IPv6，httpx 解析直接炸
for _v in ("no_proxy", "NO_PROXY"):
    _val = os.environ.get(_v, "")
    if "[" in _val:
        os.environ[_v] = ",".join(p for p in _val.split(",") if "[" not in p)

import spaces
import torch
import gradio as gr
from diffusers import QwenImage21Pipeline

# Qwen-Image-2.1：7B 视觉生成 + 大文本编码器，bf16 约 32GB，48GB ZeroGPU 可跑
pipe = QwenImage21Pipeline.from_pretrained(
    "Qwen/Qwen-Image-2.1",
    torch_dtype=torch.bfloat16,
)
pipe.to("cuda")

ASPECT_RATIOS = {
    "1:1": (1024, 1024),
    "4:3": (1200, 896),
    "3:4": (896, 1200),
    "3:2": (1264, 848),
    "2:3": (848, 1264),
    "16:9": (1376, 768),
    "9:16": (768, 1376),
}


@spaces.GPU(duration=120)
def generate(prompt, ref_image, aspect_ratio, steps, seed):
    width, height = ASPECT_RATIOS[aspect_ratio]
    generator = None
    if seed is not None and int(seed) >= 0:
        generator = torch.Generator(device="cuda").manual_seed(int(seed))
    kwargs = dict(
        prompt=prompt,
        num_inference_steps=int(steps),
        generator=generator,
    )
    if ref_image is not None:
        # 有参考图：按提示词改图，尺寸跟随原图
        kwargs["image"] = ref_image
    else:
        # 纯文生图
        kwargs["width"] = width
        kwargs["height"] = height
    image = pipe(**kwargs).images[0]
    return image


with gr.Blocks(title="主人 Qwen-Image-2.1 画坊") as demo:
    gr.Markdown(
        "# 🎨 主人 Qwen-Image-2.1 画坊\n"
        "Qwen-Image-2.1 开源权重（7B），中文提示词、中文文字渲染效果极佳。\n"
        "跑在 Hugging Face ZeroGPU 免费档上。"
    )
    prompt = gr.Textbox(label="画面描述（中文英文都行）",
                        placeholder="月下喝茶的汉服小猫，水墨画风\n（上传参考图后，这里写改图指令，如：把背景换成雪山）",
                        lines=3)
    ref_image = gr.Image(label="参考图（可选，不传就是纯文生图；传了就按上面的描述改这张图）",
                         type="pil")
    with gr.Row():
        aspect = gr.Dropdown(list(ASPECT_RATIOS.keys()), value="1:1", label="画幅")
        steps = gr.Slider(10, 50, value=30, step=1, label="步数（越多越细，也越慢）")
        seed = gr.Number(value=-1, label="种子（-1=随机）", precision=0)
    btn = gr.Button("画！", variant="primary")
    out = gr.Image(label="成品")
    btn.click(generate, inputs=[prompt, ref_image, aspect, steps, seed], outputs=out)

demo.launch()