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

device = "cpu"

print("Loading CPU model... this may take 20–40 seconds.")

pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    safety_checker=None,
    torch_dtype=torch.float32
).to(device)

def generate(img, prompt, strength, steps, guidance):
    if img is None:
        return "لطفاً عکس ورودی بده."

    result = pipe(
        prompt=prompt,
        image=img,
        strength=strength,
        num_inference_steps=steps,
        guidance_scale=guidance
    ).images[0]

    return result

with gr.Blocks(title="CPU Portrait Generator") as demo:

    gr.Markdown("## **نسخه CPU – بدون GPU – سازگار با Space رایگان**")

    with gr.Row():

        with gr.Column():
            input_img = gr.Image(type="pil", label="عکس ورودی")
            prompt = gr.Textbox(label="پرامپت")
            strength = gr.Slider(0.1, 1.0, value=0.6, label="قدرت ادیت")
            steps = gr.Slider(10, 40, value=25, label="Steps")
            guidance = gr.Slider(1, 12, value=7.5, label="Guidance Scale")
            btn = gr.Button("ساخت تصویر")

        with gr.Column():
            output = gr.Image(label="خروجی نهایی")

    btn.click(
        fn=generate,
        inputs=[input_img, prompt, strength, steps, guidance],
        outputs=output
    )

demo.launch()