import gradio as gr
import torch
from diffusers import StableDiffusionImg2ImgPipeline
from PIL import Image

torch.set_num_threads(2)

MODEL_ID = "runwayml/stable-diffusion-v1-5"

pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,
    safety_checker=None
)

pipe = pipe.to("cpu")
pipe.enable_attention_slicing()

def generate(image, prompt, strength):
    image = image.convert("RGB")
    image = image.resize((384, 384))  # CPU için daha stabil

    result = pipe(
        prompt=prompt,
        image=image,
        strength=strength,
        num_inference_steps=20,   # SD için ideal düşük kalite dengesi
        guidance_scale=7.5
    ).images[0]

    return result

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Image(type="pil", label="Upload Room Image"),
        gr.Textbox(label="Interior Prompt"),
        gr.Slider(0.4, 0.8, value=0.6, label="Strength")
    ],
    outputs=gr.Image(label="Redesigned Interior"),
    title="AI Interior Redesign (Stable Diffusion v1.5)",
)

demo.launch()