app.py

import gradio as gr
import torch
from diffusers import DiffusionPipeline, QwenImageEditPipeline
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
from qwen_vl_utils import QwenVLTokenizer
from PIL import Image

# Global variables for each model pipeline
gen_pipe = None
edit_pipe = None

# --- Model Loading Functions ---

def setup_generation_model():
    """
    Loads and quantizes the Qwen/Qwen-Image model for text-to-image generation using bitsandbytes.
    """
    global gen_pipe
    if gen_pipe is not None:
        return "Generation Model already loaded. ✨"

    model_id = "Qwen/Qwen-Image"
    device = "cuda" if torch.cuda.is_available() else "cpu"

    print(f"Loading Qwen-Image Generation Model on {device} with bitsandbytes quantization...")

    try:
        # Define BitsAndBytesConfig for 4-bit quantization
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )

        gen_model = AutoModelForCausalLM.from_pretrained(
            model_id,
            quantization_config=bnb_config if device == "cuda" else None,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True
        )
        
        gen_tokenizer = QwenVLTokenizer.from_pretrained(model_id, trust_remote_code=True)
        
        gen_pipe = DiffusionPipeline.from_pretrained(
            model_id,
            model=gen_model,
            tokenizer=gen_tokenizer,
            torch_dtype=torch.bfloat16,
            use_safetensors=True,
            trust_remote_code=True
        )
        gen_pipe.to(device)
        print("Qwen-Image Generation Model loaded and quantized successfully.")
        return "Generation Model loaded! 🚀"
    except Exception as e:
        gen_pipe = None
        return f"Generation Model setup failed. Error: {e}"

def setup_editing_model():
    """
    Loads the Qwen/Qwen-Image-Edit pipeline for image-to-image editing using bitsandbytes.
    """
    global edit_pipe
    if edit_pipe is not None:
        return "Editing Model already loaded. ✨"

    model_id = "Qwen/Qwen-Image-Edit"
    device = "cuda" if torch.cuda.is_available() else "cpu"

    print(f"Loading Qwen-Image-Edit Model on {device} with bitsandbytes quantization...")

    try:
        # Note: QwenImageEditPipeline is a custom pipeline, so direct bnb quantization might not work
        # as seamlessly as with the CausalLM model. We'll rely on it internally.
        edit_pipe = QwenImageEditPipeline.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16,
            use_safetensors=True,
            trust_remote_code=True
        )
        edit_pipe.to(device)
        print("Qwen-Image-Edit model loaded successfully.")
        return "Editing Model loaded! ✂️"
    except Exception as e:
        edit_pipe = None
        return f"Editing Model setup failed. Error: {e}"

# --- Generation and Editing Functions (remain the same as before) ---

def generate_image(prompt, negative_prompt, num_inference_steps, guidance_scale, seed):
    global gen_pipe
    if gen_pipe is None: return None, "Model not loaded.", ""
    generator = torch.Generator(device=gen_pipe.device).manual_seed(seed) if seed != -1 else None
    try:
        image = gen_pipe(prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator).images[0]
        return image, "Image generated successfully!", ""
    except Exception as e:
        return None, "An error occurred during image generation.", f"Error: {e}"

def edit_image(input_image_pil, prompt, negative_prompt, num_inference_steps, guidance_scale, true_cfg_scale, denoising_strength, seed):
    global edit_pipe
    if edit_pipe is None: return None, "Model not loaded.", ""
    if input_image_pil is None: return None, "Please upload an image.", ""
    generator = torch.Generator(device=edit_pipe.device).manual_seed(seed) if seed != -1 else None
    try:
        edited_image = edit_pipe(image=input_image_pil.convert("RGB"), prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, true_cfg_scale=true_cfg_scale, denoising_strength=denoising_strength, generator=generator).images[0]
        return edited_image, "Image edited successfully!", ""
    except Exception as e:
        return None, "An error occurred during image editing.", f"Error: {e}"

# --- Gradio UI ---

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🎨 Qwen Image Studio: Generation & Editing")
    gr.Markdown("Explore the power of Qwen models for advanced image generation and detailed editing.")

    with gr.Tab("Image Generation (Qwen/Qwen-Image)"):
        gr.Markdown("### Text-to-Image Generation")
        gr.Markdown("Create new images from text prompts. ")
        
        with gr.Row():
            gen_model_status = gr.Textbox(value="Generation Model not loaded. Click 'Load' to begin.", interactive=False, label="Model Status")
            load_gen_button = gr.Button("Load Generation Model", variant="primary")
            load_gen_button.click(fn=setup_generation_model, outputs=gen_model_status)

        with gr.Column():
            gen_prompt = gr.Textbox(label="Prompt", placeholder="A majestic dragon flying over a futuristic city at sunset, highly detailed, photorealistic", lines=2)
            gen_negative_prompt = gr.Textbox(label="Negative Prompt (Optional)", placeholder="blurry, low quality, distorted, bad anatomy", lines=1)
            
            with gr.Accordion("Advanced Settings", open=False):
                with gr.Row():
                    gen_num_steps = gr.Slider(minimum=10, maximum=150, step=1, value=50, label="Inference Steps")
                    gen_guidance_scale = gr.Slider(minimum=1.0, maximum=20.0, step=0.5, value=7.5, label="Guidance Scale")
                with gr.Row():
                    gen_seed = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)
            
            generate_button = gr.Button("Generate Image", variant="secondary")
            
            gen_output_image = gr.Image(label="Generated Image")
            gen_status_text = gr.Textbox(label="Status", interactive=False)
            gen_error_text = gr.Textbox(label="Error Details", interactive=False, visible=False)

            generate_button.click(
                fn=generate_image,
                inputs=[gen_prompt, gen_negative_prompt, gen_num_steps, gen_guidance_scale, gen_seed],
                outputs=[gen_output_image, gen_status_text, gen_error_text]
            )

    with gr.Tab("Image Editing (Qwen/Qwen-Image-Edit)"):
        gr.Markdown("### Image-to-Image Editing")
        gr.Markdown("Upload an image and provide a text prompt to transform it. This model excels at semantic and appearance editing.")

        with gr.Row():
            edit_model_status = gr.Textbox(value="Editing Model not loaded. Click 'Load' to begin.", interactive=False, label="Model Status")
            load_edit_button = gr.Button("Load Editing Model", variant="primary")
            load_edit_button.click(fn=setup_editing_model, outputs=edit_model_status)

        with gr.Column():
            edit_input_image = gr.Image(label="Upload Image to Edit", type="pil")
            edit_prompt = gr.Textbox(label="Edit Prompt", placeholder="Change the dog's fur to a vibrant blue and add a red collar", lines=2)
            edit_negative_prompt = gr.Textbox(label="Negative Prompt (Optional)", placeholder="blurry, low quality, distorted, messy", lines=1)

            with gr.Accordion("Advanced Settings", open=False):
                with gr.Row():
                    edit_num_steps = gr.Slider(minimum=10, maximum=150, step=1, value=50, label="Inference Steps")
                    edit_guidance_scale = gr.Slider(minimum=1.0, maximum=20.0, step=0.5, value=7.5, label="Guidance Scale")
                with gr.Row():
                    edit_true_cfg_scale = gr.Slider(minimum=1.0, maximum=10.0, step=0.1, value=4.0, label="True CFG Scale (for more precise control)")
                    edit_denoising_strength = gr.Slider(minimum=0.0, maximum=1.0, step=0.01, value=0.8, label="Denoising Strength (how much to change original)")
                with gr.Row():
                    edit_seed = gr.Number(label="Seed (-1 for random)", value=-1, precision=0)

            edit_button = gr.Button("Edit Image", variant="secondary")
            
            edit_output_image = gr.Image(label="Edited Image")
            edit_status_text = gr.Textbox(label="Status", interactive=False)
            edit_error_text = gr.Textbox(label="Error Details", interactive=False, visible=False)

            edit_button.click(
                fn=edit_image,
                inputs=[edit_input_image, edit_prompt, edit_negative_prompt, edit_num_steps, edit_guidance_scale, edit_true_cfg_scale, edit_denoising_strength, edit_seed],
                outputs=[edit_output_image, edit_status_text, edit_error_text]
            )

# Launch the app
demo.launch(inbrowser=True, share=False)