import spaces
import gradio as gr
import torch
from PIL import Image
from diffusers import DiffusionPipeline
import random
import uuid
import numpy as np
import time
import os

# Description for the app
DESCRIPTION = """
# Qwen Image Upscaler

Upload a low-quality or small image, and this app will use the Qwen-Image model to generate a higher-resolution, more detailed version.
"""

# --- Helper functions ---

def save_image(img: Image.Image) -> str:
    """Saves an image to a unique filename and returns the path."""
    unique_name = str(uuid.uuid4()) + ".png"
    img.save(unique_name)
    return unique_name

MAX_SEED = np.iinfo(np.int32).max

# --- Load the Qwen/Qwen-Image pipeline ---
# This single pipeline is used for both text-to-image and image-to-image (upscaling)

print("Loading Qwen-Image model...")
dtype = torch.bfloat16
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

pipe_qwen = DiffusionPipeline.from_pretrained(
    "Qwen/Qwen-Image", 
    torch_dtype=dtype
).to(device)
print("Model loaded successfully.")


# --- The main upscaler function ---

@spaces.GPU(duration=120)
def upscale_image(
    image: Image.Image,
    prompt: str,
    negative_prompt: str,
    seed: int,
    guidance_scale: float,
    randomize_seed: bool,
    num_inference_steps: int,
    progress=gr.Progress(track_tqdm=True)
):
    """
    Takes a low-resolution image and upscales it using the Qwen-Image model.
    """
    if image is None:
        raise gr.Error("No image uploaded. Please upload an image to upscale.")

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    generator = torch.Generator(device).manual_seed(seed)

    start_time = time.time()
    
    # The pipeline automatically handles upscaling when an `image` argument is provided.
    upscaled_image = pipe_qwen(
        prompt=prompt,
        negative_prompt=negative_prompt,
        image=image,  # Providing the input image triggers the upscaling/img2img mode
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        generator=generator,
        output_type="pil",
    ).images[0]

    end_time = time.time()
    duration = end_time - start_time
    
    image_path = save_image(upscaled_image)

    print(f"Upscaling finished in {duration:.2f} seconds. Seed used: {seed}")

    return image_path, seed, f"{duration:.2f}"


# --- Gradio User Interface ---

css = '''
.gradio-container {
    max-width: 840px !important;
    margin: 0 auto !important;
}
h1 {
    text-align: center;
}
footer {
    visibility: hidden;
}
'''

with gr.Blocks(css=css, theme="bethecloud/storj_theme") as demo:
    gr.Markdown(DESCRIPTION)
    
    with gr.Row():
        with gr.Column(scale=1):
            image_upload = gr.Image(
                label="Upload Low-Resolution Image", 
                type="pil",
                tool='editor'
            )
            prompt = gr.Textbox(
                label="Prompt", 
                value="ultra-detailed, high quality, 4k, 8k, masterpiece",
                placeholder="Describe the desired result (e.g., 'photorealistic, sharp focus')."
            )
            upscale_button = gr.Button("Upscale Image", variant="primary")

        with gr.Column(scale=1):
            upscaled_image_result = gr.Image(label="Upscaled Image")
            
    with gr.Accordion("Upscaler Options", open=False):
        negative_prompt = gr.Text(
            label="Negative Prompt",
            max_lines=1,
            placeholder="Enter concepts to avoid (e.g., 'blurry, pixelated').",
            value="blurry, low resolution, text, watermark, jpeg artifacts, compression",
        )
        seed = gr.Slider(
            label="Seed",
            minimum=0,
            maximum=MAX_SEED,
            step=1,
            value=0,
        )
        randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
        guidance_scale = gr.Slider(
            label="Guidance Scale",
            minimum=0.0,
            maximum=20.0,
            step=0.1,
            value=4.0,
        )
        num_inference_steps = gr.Slider(
            label="Number of Inference Steps",
            minimum=1,
            maximum=100,
            step=1,
            value=25, # Upscaling often requires fewer steps than generation from scratch
        )

    with gr.Accordion("Output Information", open=True):
         with gr.Row():
            seed_display = gr.Textbox(label="Seed used", interactive=False)
            generation_time = gr.Textbox(label="Generation time (seconds)", interactive=False)
            
    # Connect the button to the function
    upscale_button.click(
        fn=upscale_image,
        inputs=[
            image_upload,
            prompt,
            negative_prompt,
            seed,
            guidance_scale,
            randomize_seed,
            num_inference_steps
        ],
        outputs=[
            upscaled_image_result,
            seed_display,
            generation_time,
        ],
        api_name="upscale"
    )

if __name__ == "__main__":
    demo.queue(max_size=20).launch(share=False, debug=True, show_error=True)