app.py

import torch
import gradio as gr
from PIL import Image
import qrcode
from pathlib import Path
from multiprocessing import cpu_count
import requests
import io
import os
from PIL import Image
#import spaces
import numpy as np
import cv2
from pyzxing import BarCodeReader
from PIL import ImageOps
from huggingface_hub import hf_hub_download, snapshot_download
from PIL import ImageEnhance

from diffusers import (
    StableDiffusionPipeline,
    StableDiffusionControlNetImg2ImgPipeline,
    StableDiffusionControlNetPipeline,
    ControlNetModel,
    DDIMScheduler,
    DPMSolverMultistepScheduler,
    DEISMultistepScheduler,
    HeunDiscreteScheduler,
    EulerDiscreteScheduler,
)



qrcode_generator = qrcode.QRCode(
    version=1,
    error_correction=qrcode.ERROR_CORRECT_H,
    box_size=10,
    border=4,
)


# Define available models
CONTROLNET_MODELS = {
    "QR Code Monster": "monster-labs/control_v1p_sd15_qrcode_monster",
    "QR Code": "DionTimmer/controlnet_qrcode-control_v1p_sd15",
    # Add more ControlNet models here
}

DIFFUSION_MODELS = {
    "GhostMix": "sinkinai/GhostMix-V2-BakedVae",
    # Add more diffusion models here
}

# Global variables to store loaded models
loaded_controlnet = None
loaded_pipe = None

def load_models_on_launch():
    global loaded_controlnet, loaded_pipe
    print("Loading models on launch...")
    controlnet_path = snapshot_download(CONTROLNET_MODELS["QR Code Monster"])
    loaded_controlnet = ControlNetModel.from_pretrained(
        controlnet_path,
        torch_dtype=torch.float16
    ).to("mps")

    diffusion_path = snapshot_download(DIFFUSION_MODELS["GhostMix"])
    loaded_pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
        diffusion_path,
        controlnet=loaded_controlnet,
        torch_dtype=torch.float16,
        safety_checker=None,
    ).to("mps")
    print("Models loaded successfully!")

# Modify the load_models function to use global variables
def load_models(controlnet_model, diffusion_model):
    global loaded_controlnet, loaded_pipe
    if loaded_controlnet is None or loaded_pipe is None:
        load_models_on_launch()
    return loaded_pipe

# Add new functions for image adjustments
def adjust_image(image, brightness, contrast, saturation):
    if image is None:
        return None
    
    img = Image.fromarray(image) if isinstance(image, np.ndarray) else image
    
    if brightness != 1:
        img = ImageEnhance.Brightness(img).enhance(brightness)
    if contrast != 1:
        img = ImageEnhance.Contrast(img).enhance(contrast)
    if saturation != 1:
        img = ImageEnhance.Color(img).enhance(saturation)
    
    return np.array(img)

def resize_for_condition_image(input_image: Image.Image, resolution: int):
    input_image = input_image.convert("RGB")
    W, H = input_image.size
    k = float(resolution) / min(H, W)
    H *= k
    W *= k
    H = int(round(H / 64.0)) * 64
    W = int(round(W / 64.0)) * 64
    img = input_image.resize((W, H), resample=Image.LANCZOS)
    return img


SAMPLER_MAP = {
    "DPM++ Karras SDE": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True, algorithm_type="sde-dpmsolver++"),
    "DPM++ Karras": lambda config: DPMSolverMultistepScheduler.from_config(config, use_karras=True),
    "Heun": lambda config: HeunDiscreteScheduler.from_config(config),
    "Euler": lambda config: EulerDiscreteScheduler.from_config(config),
    "DDIM": lambda config: DDIMScheduler.from_config(config),
    "DEIS": lambda config: DEISMultistepScheduler.from_config(config),
}

def scan_qr_code(image):
    # Convert gradio image to PIL Image if necessary
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    
    # Convert to grayscale
    gray_image = image.convert('L')
    
    # Convert to numpy array
    np_image = np.array(gray_image)
    
    # Method 1: Using qrcode library
    try:
        qr = qrcode.QRCode()
        qr.add_data('')
        qr.decode(gray_image)
        return qr.data.decode('utf-8')
    except Exception:
        pass
    
    # Method 2: Using OpenCV
    try:
        qr_detector = cv2.QRCodeDetector()
        retval, decoded_info, points, straight_qrcode = qr_detector.detectAndDecodeMulti(np_image)
        if retval:
            return decoded_info[0]
    except Exception:
        pass
    
    # Method 3: Fallback to zxing-cpp
    try:
        reader = BarCodeReader()
        results = reader.decode(np_image)
        if results:
            return results[0].parsed
    except Exception:
        pass
    
    return None

def invert_image(image):
    if image is None:
        return None
    if isinstance(image, np.ndarray):
        return 255 - image
    elif isinstance(image, Image.Image):
        return ImageOps.invert(image.convert('RGB'))
    else:
        raise ValueError("Unsupported image type")

def invert_displayed_image(image):
    if image is None:
        return None
    inverted = invert_image(image)
    if isinstance(inverted, np.ndarray):
        return Image.fromarray(inverted)
    return inverted

#@spaces.GPU()
def inference(
    qr_code_content: str,
    prompt: str,
    negative_prompt: str,
    guidance_scale: float = 10.0,
    controlnet_conditioning_scale: float = 2.0,
    strength: float = 0.8,
    seed: int = -1,
    init_image: Image.Image | None = None,
    use_qr_code_as_init_image = True,
    sampler = "DPM++ Karras SDE",
    bg_color: str = "white",
    qr_color: str = "black",
    invert_final_image: bool = False,
    invert_init_image: bool = False,
    controlnet_model: str = "QR Code Monster",
    diffusion_model: str = "GhostMix",
):
    try:
        progress = gr.Progress()
        # Load models based on user selection
        progress(0, desc="Downloading models...")
        pipe = load_models(controlnet_model, diffusion_model)
        progress(0.5, desc="Models downloaded, preparing for inference...")
        
        
        if prompt is None or prompt == "":
            raise gr.Error("Prompt is required")

        if qr_code_content == "":
            raise gr.Error("QR Code Content is required")

        pipe.scheduler = SAMPLER_MAP[sampler](pipe.scheduler.config)

        if seed == -1:
            seed = torch.randint(0, 2**32 - 1, (1,)).item()
        generator = torch.manual_seed(seed)

        print("Generating QR Code from content")
        qr = qrcode.QRCode(
            version=1,
            error_correction=qrcode.constants.ERROR_CORRECT_H,
            box_size=10,
            border=4,
        )
        qr.add_data(qr_code_content)
        qr.make(fit=True)

        qrcode_image = qr.make_image(fill_color=qr_color, back_color=bg_color)
        qrcode_image = resize_for_condition_image(qrcode_image, 768)

        # Determine which image to use as init_image and control_image
        if use_qr_code_as_init_image:
            init_image = qrcode_image
            control_image = qrcode_image
        else:
            control_image = qrcode_image
            if init_image is None:
                # If no init_image provided, set strength to 1.0 to generate a new image
                strength = 1.0

        # Adjust strength if using an init_image
        if init_image is not None:
            strength = min(strength, 0.8)  # Cap strength at 0.8 when using init_image

        # Invert init_image if requested
        if invert_init_image and init_image is not None:
            init_image = invert_image(init_image)

        out = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt,
            image=init_image,
            control_image=control_image,
            width=1024,
            height=1024,
            guidance_scale=float(guidance_scale),
            controlnet_conditioning_scale=float(controlnet_conditioning_scale),
            generator=generator,
            strength=float(strength),
            num_inference_steps=100,
        )
        
        final_image = out.images[0]
        
        if invert_final_image:
            final_image = invert_image(final_image)
        
        return final_image, seed
    except Exception as e:
        print(f"Error in inference: {str(e)}")
        return Image.new('RGB', (1024, 1024), color='white'), -1

def invert_init_image_display(image):
    if image is None:
        return None
    inverted = invert_image(image)
    if isinstance(inverted, np.ndarray):
        return Image.fromarray(inverted)
    return inverted

with gr.Blocks(theme='Hev832/Applio') as blocks:
    gr.Markdown(
        """
        ![Yamamoto Logo](https://cdn-uploads.huggingface.co/production/uploads/64740cf7485a7c8e1bd51ac9/_VyYxp5qE_nRZ_LJqBxmL.webp)
        # 🎨 Yamamoto QR Code Art Generator
        Transform Your QR Codes into Brand Masterpieces
        """
    )

    with gr.Tabs():
        with gr.TabItem("1. Input & Design"):
            with gr.Row():
                with gr.Column(scale=1):
                    qr_code_content = gr.Textbox(
                        label="QR Code Content",
                        placeholder="Enter URL or text for your QR code",
                        info="This is what your QR code will link to or display when scanned.",
                    )
                    prompt = gr.Textbox(
                        label="Artistic Prompt",
                        placeholder="Describe the style or theme for your QR code art",
                        value="A high-resolution, photo-realistic minimalist rendering of Mount Fuji, depicted as a sharp, semi-realistic silhouette of a mountain range on the horizon. The mountain evokes strength and motion with clean, crisp lines and a sense of natural flow. The scene should feature detailed snow textures, subtle highlights on the mountain ridges, and a powerful yet serene atmosphere. The rendering should emphasize the strength of the mountain with a focus on clarity and precision in both texture and light. (Sharp outlines:1.5), (Photo-realistic:1.4), (Detailed textures:1.3), (Minimalist:1.3), (Semi-realistic:1.3), (Monochrome contrast:1.2), (Crisp detail:1.2), (Evoking strength:1.2), inspired by traditional Japanese woodblock prints, nature photography, and minimalist design principles.",
                        info="Be specific and creative! This guides the AI in creating your unique QR code art.",
                    )
                    negative_prompt = gr.Textbox(
                        label="Elements to Avoid",
                        placeholder="Describe what you don't want in the image",
                        value="ugly, disfigured, low quality, blurry, nsfw, bad_pictures, (bad_prompt_version2:0.8), EasyNegative, 3d, cartoon, anime, sketches, (worst quality:2), (low quality:2), (normal quality:2), lowres, normal quality, ((monochrome)), ((grayscale)), poorly drawn, distorted, overexposed, flat shading, bad proportions, deformed, pixelated, messy details, lack of contrast, unrealistic textures, bad anatomy, rough edges, low resolution, text artifacts.",
                        info="List elements or styles you want to avoid in your QR code art.",
                    )
                
                with gr.Column(scale=1):
                    with gr.Accordion("QR Code Customization", open=True):
                        bg_color = gr.ColorPicker(
                            label="Background Color", 
                            value="#FFFFFF",
                            info="Choose the background color for the QR code"
                        )
                        qr_color = gr.ColorPicker(
                            label="QR Code Color", 
                            value="#000000",
                            info="Choose the color for the QR code pattern"
                        )
                    
                    with gr.Accordion("Use Your Own Image as a Reference", open=False):
                        init_image = gr.Image(label="Reference Image", type="pil")
                        use_qr_code_as_init_image = gr.Checkbox(
                            label="Use QR code as base image", 
                            value=True, 
                            interactive=True,
                            info="Uncheck to use your own image for generation"
                        )
                        invert_init_image_button = gr.Button("Invert Init Image")

        with gr.TabItem("2. Advanced Settings"):
            with gr.Row():
                with gr.Column(scale=1):
                    with gr.Accordion("Art Generation Controls", open=True):
                        controlnet_conditioning_scale = gr.Slider(
                            minimum=0.0, maximum=5.0, step=0.01, value=2,
                            label="QR Code Visibility in Image",
                        )
                        strength = gr.Slider(
                            minimum=0.0, maximum=1.0, step=0.01, value=0.9,
                            label="Artistic Freedom for the AI",
                        )
                        guidance_scale = gr.Slider(
                            minimum=0.0, maximum=50.0, step=0.25, value=7.5,
                            label="How closely the AI follows the Prompt",
                        )
                    
                    with gr.Accordion("Model Selection", open=True):
                        controlnet_model_dropdown = gr.Dropdown(
                            choices=list(CONTROLNET_MODELS.keys()),
                            value="QR Code Monster",
                            label="ControlNet Model",
                            info="Select the ControlNet model for QR code generation"
                        )
                        diffusion_model_dropdown = gr.Dropdown(
                            choices=list(DIFFUSION_MODELS.keys()),
                            value="GhostMix",
                            label="Diffusion Model",
                            info="Select the main diffusion model for image generation"
                        )

                with gr.Column(scale=1):
                    with gr.Accordion("Generation Settings", open=True):
                        sampler = gr.Dropdown(
                            choices=list(SAMPLER_MAP.keys()), 
                            value="DPM++ Karras SDE", 
                            label="Art Style the AI uses to create the image",
                        )
                        seed = gr.Slider(
                            minimum=-1, maximum=9999999999, step=1, value=-1,
                            label="Creative Seed for the Image Generation",
                            randomize=False,                    
                        )
                    
                    with gr.Accordion("Additional Options", open=False):
                        invert_final_image = gr.Checkbox(
                            label="Invert Final Image", 
                            value=False,
                            info="Check this to invert the colors of the final image",
                        )

        with gr.TabItem("3. Generate & Refine"):
            with gr.Row():
                with gr.Column(scale=1):
                    run_btn = gr.Button("🎨 Create Your QR Art", variant="primary")
                    result_image = gr.Image(label="Your Artistic QR Code")
                    used_seed = gr.Number(label="Seed Used", interactive=False)
                
                with gr.Column(scale=1):
                    with gr.Accordion("Image Adjustment", open=True):
                        brightness = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Brightness")
                        contrast = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Contrast")
                        saturation = gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="Saturation")
                        invert_button = gr.Button("Invert Image")
                    
                    with gr.Accordion("QR Code Verification", open=True):
                        scan_button = gr.Button("Verify QR Code Works")
                        scan_result = gr.Textbox(label="Validation Result of QR Code", interactive=False)

    gr.Markdown(
        """
        ### 🔍 Tips for Optimizing Your QR Art
        - If the QR code isn't scannable, try adjusting Brightness, Contrast, and Saturation.
        - For more artistic flair, increase 'Artistic Freedom' in Advanced Settings.
        - To make the QR code clearer, raise 'QR Code Visibility' in Advanced Settings.
        - Experiment with different prompts and settings to find your perfect style!
        """
    )

    def scan_and_display(image):
        if image is None:
            return "No image to scan"
        
        scanned_text = scan_qr_code(image)
        if scanned_text:
            return f"Scanned successfully: {scanned_text}"
        else:
            return "Failed to scan QR code. Try adjusting the settings for better visibility."

    def invert_displayed_image(image):
        if image is None:
            return None
        return invert_image(image)

    scan_button.click(
        scan_and_display,
        inputs=[result_image],
        outputs=[scan_result]
    )

    invert_button.click(
        invert_displayed_image,
        inputs=[result_image],
        outputs=[result_image]
    )

    invert_init_image_button.click(
        invert_init_image_display,
        inputs=[init_image],
        outputs=[init_image]
    )

    brightness.change(
        adjust_image,
        inputs=[result_image, brightness, contrast, saturation],
        outputs=[result_image]
    )
    contrast.change(
        adjust_image,
        inputs=[result_image, brightness, contrast, saturation],
        outputs=[result_image]
    )
    saturation.change(
        adjust_image,
        inputs=[result_image, brightness, contrast, saturation],
        outputs=[result_image]
    )

    run_btn.click(
        inference,
        inputs=[
            qr_code_content,
            prompt,
            negative_prompt,
            guidance_scale,
            controlnet_conditioning_scale,
            strength,
            seed,
            init_image,
            use_qr_code_as_init_image,
            sampler,
            bg_color,
            qr_color,
            invert_final_image,
            controlnet_model_dropdown,
            diffusion_model_dropdown,
        ],
        outputs=[result_image, used_seed],
        concurrency_limit=20
    )

# Load models on launch
load_models_on_launch()

blocks.queue(max_size=20)
blocks.launch(share=False, show_api=True)