import gradio as gr
import torch
import numpy as np
import cv2
from PIL import Image

from diffusers import (
ControlNetModel,
StableDiffusionControlNetPipeline,
StableDiffusionImg2ImgPipeline,
)

device = "cuda" if torch.cuda.is_available() else "cpu"

Load ControlNet (keeps structure)

controlnet = ControlNetModel.from_pretrained(
"lllyasviel/sd-controlnet-canny",
torch_dtype=torch.float32
)

pipe_control = StableDiffusionControlNetPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
controlnet=controlnet,
torch_dtype=torch.float32,
)

pipe_control = pipe_control.to(device)

Load Img2Img (makes it realistic)

pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained(
"runwayml/stable-diffusion-v1-5",
torch_dtype=torch.float32,
)

pipe_img2img = pipe_img2img.to(device)

def process(image, prompt):
# Convert to edges (structure)
img = np.array(image)
edges = cv2.Canny(img, 100, 200)
edges = np.stack([edges]*3, axis=2)
edges = Image.fromarray(edges)

# First pass (keep layout)
result1 = pipe_control(
    prompt=prompt,
    image=edges,
    num_inference_steps=20
).images[0]

# Second pass (refine)
result2 = pipe_img2img(
    prompt=prompt,
    image=result1,
    strength=0.35,
    num_inference_steps=20
).images[0]

return edges, result1, result2

with gr.Blocks() as demo:
gr.Markdown("## Simple 2-Step Render AI")

input_image = gr.Image(type="pil", label="Upload drawing")
prompt = gr.Textbox(value="architectural render, realistic materials, glass curtain wall, brick facade")

btn = gr.Button("Generate")

out1 = gr.Image(label="Edges")
out2 = gr.Image(label="Step 1 Render")
out3 = gr.Image(label="Final Render")

btn.click(process, inputs=[input_image, prompt], outputs=[out1, out2, out3])

demo.launch()