# Run Stable Diffusion to generate images from text prompts.
# This script uses the Modal framework to run Stable Diffusion in a cloud environment.
# It requires the `modal` package and a Hugging Face token stored in a Modal secret.
# Make sure to set up the Modal environment and install the necessary dependencies.

# Usage: modal run stablediffusion.py

from icecream import ic
import io
import os

import modal

app = modal.App()

@app.function(
    image=modal.Image.debian_slim().pip_install(
        "icecream",
        "torch",
        "diffusers[torch]",
        "transformers",
        "ftfy"
    ),
    secrets=[modal.Secret.from_name("huggingface-secret")],
    gpu="any",
)
def run_stable_diffusion(prompt: str):
    from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import StableDiffusionPipeline

    pipe = StableDiffusionPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5",
        use_auth_token=os.environ["HF_TOKEN"],
    ).to("cuda")

    image = pipe(prompt, num_inference_steps=10).images[0] # type: ignore

    buf = io.BytesIO()
    image.save(buf, format="PNG")
    img_bytes = buf.getvalue()

    return img_bytes


@app.local_entrypoint()
def main():
    prompt = "Wu-Tang Clan climbing Mount Everest"
    # prompt = "A robot dog walking down a vineyard"  # Example prompt

    # out_path = "/tmp/output.png"
    out_path = "stablediffusion_output.png"

    # ic(os.getcwd())
    # img_bytes = b"<image_bytes>"  # Placeholder for the actual image bytes

    print("DEBUG: Starting Stable Diffusion with prompt:", prompt)
    img_bytes = run_stable_diffusion.remote(prompt=prompt)
    print("DEBUG: Writing img_bytes length:", len(img_bytes))
    with open(out_path, "wb") as f:
        f.write(img_bytes)
    print("DEBUG: Image saved to ", out_path)


if __name__ == "__main__":
    main()
    print("Image saved to /tmp/output.png")
    print("Run `modal deploy` to deploy this app.")
    print("Run `modal serve` to serve this app locally.")
    print("Run `modal run` to run this app in the cloud.")
    print("Run `modal logs` to view the logs of this app.")
    print("Run `modal shell` to open a shell in the cloud environment.")
    print("Run `modal run --help` to see all available options.")
    print("Run `modal deploy --help` to see all available options for deployment.")