import torch
import gradio as gr
from diffusers import QwenImagePipeline, QwenImageTransformer2DModel
from transformers import Qwen2_5_VLForConditionalGeneration

# Set device and dtype
torch_dtype = torch.bfloat16

# Load models
transformer = QwenImageTransformer2DModel.from_pretrained(
    "OzzyGT/Qwen-Image-2512-bnb-4bit-transformer",
    torch_dtype=torch_dtype,
    device_map="cpu"  # offloading to CPU
)
text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
    "OzzyGT/Qwen-Image-2512-bnb-4bit-text-encoder",
    torch_dtype=torch_dtype,
    device_map="cpu"
)

pipe = QwenImagePipeline.from_pretrained(
    "Qwen/Qwen-Image-2512",
    transformer=transformer,
    text_encoder=text_encoder,
    torch_dtype=torch_dtype
)

pipe.enable_model_cpu_offload()

# Define the Gradio function
def generate_image(prompt, negative_prompt, width=1664, height=928, seed=42):
    generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        num_inference_steps=28,
        true_cfg_scale=4.0,
        generator=generator,
    ).images[0]
    return image

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## Qwen Image Generation")
    
    with gr.Row():
        with gr.Column():
            prompt_input = gr.Textbox(label="Prompt", lines=8, placeholder="Enter your prompt here...")
            negative_prompt_input = gr.Textbox(label="Negative Prompt", lines=4, placeholder="Enter negative prompt...")
            width_input = gr.Slider(256, 2048, value=1664, step=64, label="Width")
            height_input = gr.Slider(256, 2048, value=928, step=64, label="Height")
            seed_input = gr.Number(value=42, label="Random Seed")
            generate_btn = gr.Button("Generate Image")
        with gr.Column():
            output_image = gr.Image(label="Generated Image")
    
    generate_btn.click(
        fn=generate_image,
        inputs=[prompt_input, negative_prompt_input, width_input, height_input, seed_input],
        outputs=output_image
    )

demo.launch()