import gradio as gr
import torch
from PIL import Image
from diffusers import QwenImageEditPlusPipeline

MODEL_ID = "Qwen/Qwen-Image-Edit-2509"
LORA_REPO = "lovis93/next-scene-qwen-image-lora-2509"
LORA_FILE = "next-scene_lora_v1-3000.safetensors"

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

pipe = QwenImageEditPlusPipeline.from_pretrained(MODEL_ID, torch_dtype=dtype).to(device)
pipe.load_lora_weights(LORA_REPO, weight_name=LORA_FILE)

def next_scene(image, prompt, steps, true_cfg_scale, lora_strength, seed):
    gen = None
    if seed and int(seed) != 0:
        gen = torch.Generator(device=device).manual_seed(int(seed))
    try:
        pipe.set_adapters(["default"], adapter_weights=[float(lora_strength)])
    except Exception:
        pass

    kwargs = dict(
        image=[image],
        prompt=prompt,
        num_inference_steps=int(steps),
        guidance_scale=1.0,
        generator=gen,
    )
    try:
        kwargs["true_cfg_scale"] = float(true_cfg_scale)
    except Exception:
        pass

    out = pipe(**kwargs)
    return out.images[0]

with gr.Blocks() as demo:
    gr.Markdown("## Next Scene — Qwen-Image-Edit-2509 + LoRA")
    with gr.Row():
        with gr.Column():
            inp_img = gr.Image(type="pil", label="Входной кадр (старт сцены)")
            prompt = gr.Textbox(
                label='Промпт (начинайте с "Next Scene: ...")',
                value='Next Scene: camera pulls back revealing the riverside at sunset, soft rim light, subtle lens flare.'
            )
            steps = gr.Slider(4, 60, value=40, step=1, label="Steps")
            true_cfg = gr.Slider(1.0, 6.0, value=3.0, step=0.5, label="true_cfg_scale")
            lora_strength = gr.Slider(0.0, 1.2, value=0.75, step=0.05, label="LoRA strength")
            seed = gr.Number(value=0, label="Seed (0 = random)")
            btn = gr.Button("Сгенерировать следующий кадр")
        with gr.Column():
            out_img = gr.Image(label="Результат")

    btn.click(next_scene, [inp_img, prompt, steps, true_cfg, lora_strength, seed], [out_img])

if __name__ == "__main__":
    demo.launch()