import gradio as gr
import numpy as np
from PIL import Image
import openvino_genai as ov_genai
import subprocess, sys
import os
subprocess.run([sys.executable, "download_models.py"], check=True)

# Core count

BACKEND = os.environ.get("BACKEND", "genai")  # "optimum" or "genai"

MODEL_PATH = "LCM_Dreamshaper_v7-int8-ov"
DEVICE = "CPU"



if BACKEND == "optimum":
    from optimum.intel import OVDiffusionPipeline

    print("Loading with Optimum-Intel...")
    pipe =  OVDiffusionPipeline.from_pretrained(
        MODEL_PATH,
        device=DEVICE,
        safety_checker=None
    )

    def generate(prompt, negative_prompt, num_steps, guidance_scale, seed):
        generator = np.random.RandomState(seed if seed != -1 else None)
        result = pipe(
            prompt=prompt,
            negative_prompt=negative_prompt or None,
            num_inference_steps=num_steps,
            guidance_scale=guidance_scale,
            generator=generator,
        )
        return result.images[0]

else:
    import openvino_genai as ov_genai
    print("Loading with OpenVINO GenAI...")
    pipe = ov_genai.Text2ImagePipeline(MODEL_PATH, DEVICE)

    def generate(prompt, negative_prompt, num_steps, guidance_scale, seed):
        actual_seed = seed if seed != -1 else np.random.randint(0, 2**31)
        image_tensor = pipe.generate(
            prompt,
            num_inference_steps=num_steps,
            guidance_scale=guidance_scale,
            rng_seed=actual_seed,
        )
        return Image.fromarray(np.array(image_tensor.data[0]))


with gr.Blocks(title="SD1.5 3d interior design – OpenVINO") as demo:
    gr.Markdown(
        """
        # 🏠 SD1.5 Interior LoRA + LCM-style Model— OpenVINO INT8
        Fast CPU inference with Latent Consistency Model.  
        4 steps is usually enough — crank it to 8 for more detail.
        """
    )

    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="photo, full body man, cinematic lighting...",
                lines=3,
            )
            negative_prompt = gr.Textbox(
                label="Negative Prompt",
                placeholder="blurry, low quality, watermark...",
                lines=2,
                value="blurry, low quality, artifacts, watermark",
            )
            with gr.Row():
                num_steps = gr.Slider(
                    label="Inference Steps",
                    minimum=1,
                    maximum=16,
                    value=4,
                    step=1,
                )
                guidance_scale = gr.Slider(
                    label="Guidance Scale",
                    minimum=0.0,
                    maximum=3.0,
                    value=0.8,
                    step=0.1,
                )
            seed = gr.Number(
                label="Seed (-1 = random)",
                value=-1,
                precision=0,
            )
            btn = gr.Button("Generate", variant="primary")

        with gr.Column(scale=1):
            output = gr.Image(label="Generated Image", type="pil")

    btn.click(
        fn=generate,
        inputs=[prompt, negative_prompt, num_steps, guidance_scale, seed],
        outputs=output,
    )

    gr.Examples(
        examples=[
            ["isometric ,3d render,interior a living room with a couch, chair, table and clock,indoors, book, pillow, no humans, window, bed, chair, table, plant, curtains, scenery, couch, wooden floor, clock, lamp, alarm clock","blurry, low quality", 4, 1.0, 42],
            ["isometric ,3d render,interior a bedroom with a bed, desk and computer monitor in a neon frame,book, pillow, no humans, window, bed, night, chair, plant, scenery, desk, lamp, computer, monitor","blurry, low quality", 8, 1.0, 7],
            ["isometric ,3d render,interior a small kitchen with a table and chairs,food, indoors, no humans, window, chair, table, bottle, scenery, plate, kitchen, frying pan, sink, rug, stove, cutting board","blurry, low quality", 8, 0.8, 123],
        ],
        inputs=[prompt, negative_prompt, num_steps, guidance_scale, seed],
        outputs=output,
        fn=generate,
        cache_examples=False,
    )

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