import torch
import gradio as gr
from diffusers import StableDiffusionControlNetImg2ImgPipeline, ControlNetModel
from diffusers.utils import load_image
import qrcode

# Set device CPU karena tidak ada GPU
device = "cpu"

# Load ControlNet model
controlnet = ControlNetModel.from_pretrained(
    "DionTimmer/controlnet_qrcode-control_v1p_sd15",
    torch_dtype=torch.float32
)

# Load pipeline
pipe = StableDiffusionControlNetImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    controlnet=controlnet,
    torch_dtype=torch.float32,
    safety_checker=None
).to(device)


def generate_qr(prompt, url_text):
    # 1. Buat QR dasar dari teks/URL
    qr = qrcode.QRCode(
        version=1,
        error_correction=qrcode.constants.ERROR_CORRECT_H,
        box_size=10,
        border=4,
    )
    qr.add_data(url_text)
    qr.make(fit=True)

    qr_img = qr.make_image(fill_color="black", back_color="white").convert("RGB")

    # 2. Ubah ke gaya artistik dengan ControlNet
    image = pipe(
        prompt,
        image=qr_img,
        num_inference_steps=20,
        guidance_scale=7.5
    ).images[0]

    return image


demo = gr.Interface(
    fn=generate_qr,
    inputs=[
        gr.Textbox(label="Prompt / Deskripsi Gaya Gambar"),
        gr.Textbox(label="Teks atau URL untuk QR"),
    ],
    outputs=gr.Image(label="QR Code Artistik")
)

if __name__ == "__main__":
    demo.launch()