from torch import autocast
import requests
import torch
from PIL import Image
from io import BytesIO

from diffusers import StableDiffusionImg2ImgPipeline

import os
MY_SECRET_TOKEN=os.environ.get('HF_TOKEN_SD')
YOUR_TOKEN=MY_SECRET_TOKEN
# load the pipeline
device = "cpu"
model_id_or_path = "CompVis/stable-diffusion-v1-4"
pipe = StableDiffusionImg2ImgPipeline.from_pretrained(
    model_id_or_path,
    revision="fp16", 
    torch_dtype=torch.float,
    use_auth_token=YOUR_TOKEN
)
# or download via git clone https://huggingface.co/CompVis/stable-diffusion-v1-4
# and pass `model_id_or_path="./stable-diffusion-v1-4"` without having to use `use_auth_token=True`.
pipe = pipe.to(device)

# let's download an initial image
url = "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/assets/stable-samples/img2img/sketch-mountains-input.jpg"

response = requests.get(url)
init_image = Image.open(BytesIO(response.content)).convert("RGB")
init_image = init_image.resize((768, 512))

prompt = "Lively, illustration of a king, portrait, fantasy, intricate, Scenic, hyperdetailed, hyper realistic king-hearthstone, unreal engine, 4k, smooth, sharp focus, intricate, cinematic lighting, highly detailed, octane, digital painting, artstation, concept art, vibrant colors, Cinema4D, WLOP, 3d render, in the style of hearthstone art by Artgerm and greg rutkowski and magali villeneuve, martina jackova, Giger"

with autocast(device): #"cuda"):
    images = pipe(prompt=prompt, init_image=init_image, strength=0.75, guidance_scale=7.5).images

images[0].save("fantasy_landscape.png")