import torch
from PIL import Image
from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline
from diffusers import StableDiffusionImg2ImgPipeline
from diffusers import EulerAncestralDiscreteScheduler

#provide your Hugging Face Authentication token
#you can obtain your token from https://huggingface.co/settings/tokens
auth_token = input("Enter your Hugging Face token: ")

###   image generation   ###

#using stable diffusion version 2.1 for image generation
modelid = "stabilityai/stable-diffusion-2-1"
device = "cuda" if torch.cuda.is_available() else "cpu"

pipe = StableDiffusionPipeline.from_pretrained(
    modelid,
    revision="fp16",
    torch_dtype= torch.float16,
    use_auth_token=auth_token,
    low_cpu_mem_usage=True
)
#using EulerAncestralDiscreteScheduler for sharper and detailed images

pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)


###    image modification    ###
#using StableDiffusionImg2ImgPipeline for image to image generation
pipe.to(device)

pipe_img2img = StableDiffusionImg2ImgPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16, low_cpu_mem_usage=True).to(device)
pipe_img2img.to(device)

import gradio as gr
# Function to generate image from text
def generate_image(prompt, guidance_scale):
    global stored_image
    image = pipe(prompt, guidance_scale=guidance_scale).images[0]
    stored_image = image
    return image

# Function to modify image (Img2Img)
def modify_image(prompt, strength=0.5):
    global stored_image
    if stored_image is None:
        return "No generated image available. Generate an image first!"

    stored_image = stored_image.convert("RGB")
    modified_image = pipe_img2img(prompt=prompt, image=stored_image, strength=strength, guidance_scale=8.5).images[0]
    return modified_image

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# Stable Diffusion - Generate & Modify Images")
    with gr.Group():
      with gr.Row():
        prompt_input = gr.Textbox(label="Enter Prompt", placeholder="A futuristic city at sunset")
      with gr.Row():
        guidance_input = gr.Slider(1.0, 15.0, value=8.5, label="Guidance Scale (More guidance meansthe model follows the prompt very strictly but is less creative)")  # User can adjust guidance

    with gr.Row():
      generate_btn = gr.Button("Generate")
    with gr.Row():
      output_image = gr.Image(label="Generated Image")

    generate_btn.click(generate_image, inputs=[prompt_input, guidance_input], outputs=output_image)


    # Modify Image (After Generation)
    with gr.Tab("Modify Image"):
        modify_prompt = gr.Textbox(label="Enter modification prompt")
        strength_slider = gr.Slider(0.1, 1.0, value=0.5, label="Strength (Higher = More Change)")
        modify_btn = gr.Button("Modify")
        modified_output = gr.Image(label="Modified Image")

        modify_btn.click(modify_image, inputs=[modify_prompt, strength_slider], outputs=modified_output)

# Launch Gradio App
demo.launch(share=True)