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

print("Loading model...")
pipe = QwenImageEditPlusPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2511",
    torch_dtype=torch.float16,
    use_safetensors=True,
    low_cpu_mem_usage=True
)
pipe.enable_sequential_cpu_offload()  # Экономит RAM на CPU Space

def edit_image(image_b64, prompt):
    image_b64 = image_b64.split(',')[1] if ',' in image_b64 else image_b64
    image_bytes = base64.b64decode(image_b64)
    image = Image.open(io.BytesIO(image_bytes)).resize((512,512)).convert('RGB')  # Resize для скорости
    result = pipe(image=image, prompt=prompt, num_inference_steps=15).images[0]
    buffer = io.BytesIO()
    result.save(buffer, format='JPEG', quality=90)
    return 'data:image/jpeg;base64,' + base64.b64encode(buffer.getvalue()).decode()

demo = gr.Interface(
    fn=edit_image,
    inputs=[gr.Textbox(lines=5, placeholder="data:image;base64,...", label="Image base64"), gr.Textbox(label="Prompt")],
    outputs=gr.Image(label="Result")
)
demo.launch(server_port=7860)