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

# Load the model
model_id_or_path = "stablediffusionapi/realistic-vision-v20-2047"
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(model_id_or_path, torch_dtype=torch.float32).to(device)

def generate_image(content_img, style_img, prompt, guidance_scale, strength):
    if content_img is None or style_img is None:
        return None

    content_image = Image.open(content_img).convert("RGB").resize((450, 450))
    style_image = Image.open(style_img).convert("RGB").resize((450, 450))

    seed_value = torch.randint(0, 4294967295, size=(1,)).item()

    result = pipe(
        image=content_image,
        guidance=style_image,
        guidance_scale=guidance_scale,
        strength=strength,
        num_inference_steps=100,
        prompt=prompt,
        seed=seed_value
    ).images[0]

    return result

with gr.Blocks() as demo:
    gr.Markdown("# Advanced Neural Style Transfer with Stable Diffusion")

    with gr.Row():
        content_img = gr.Image(type="filepath", label="Upload Content Image")
        style_img = gr.Image(type="filepath", label="Upload Style Image")


    with gr.Row():
        prompt = gr.Textbox(label="Enter Stable Diffusion Prompt", placeholder="e.g., 'Van Gogh painting'")
        generate_btn = gr.Button("Generate")

    with gr.Row():
        guidance_scale = gr.Slider(5.0, 20.0, value=16.0, label="Guidance Scale: #light (content image emphasized) 5.0 - 10.0 #haevy (style image emphasized) 15.0 - 20.0")
        strength = gr.Slider(0.3, 1.0, value=0.45, label="Strength")

    output_image = gr.Image(label="Generated Image")

    generate_btn.click(generate_image, inputs=[content_img, style_img, prompt, guidance_scale, strength], outputs=output_image)

demo.launch()