#sshfs philippe@station-ubuntu-c34:/home/philippe/uabox/2022-2023/sandbox/stable-diffusion ~/remote

import random
import torch
from diffusers import StableDiffusionPipeline, EulerDiscreteScheduler
#from diffusers.utils import is_accelerate_available
#from diffusers import StableDiffusionPipeline
#from transformers import CLIPFeatureExtractor, CLIPVisionModelWithProjection

def init_gpu():
    model_id = "stabilityai/stable-diffusion-2"

    # Use the Euler scheduler here instead
    scheduler = EulerDiscreteScheduler.from_pretrained(model_id, subfolder="scheduler")
    pipe = StableDiffusionPipeline.from_pretrained(model_id, scheduler=scheduler, torch_dtype=torch.float16)
    pipe = pipe.to("cuda")


# prompt = "Che Guevara en Ferrari"

# negPrompt = ""
# # 600 x 400
# # 720 x 1280
# width = 600
# height = 800
# guidance_scale = 7.5

def generate_image(prompt, negPrompt, width, height, guidance_scale, seed, sampling_steps, batch_count, utilisateur):
    if seed == -1:
        root_seed = random.randint(0,4294967295)
    else:
        root_seed = seed

    for i in range(1,batch_count):
        seed = i
        
        print (i)
        image = pipe(prompt,
                    negative_prompt = negPrompt,
                    guidance_scale = guidance_scale,
                    width=width,
                    height=height, 
                    num_inference_steps=sampling_steps,
                    #lms = LMSDiscreteScheduler( beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear" ),
                    generator=torch.Generator('cuda').manual_seed(root_seed+seed)
                    #generator=torch.Generator('cuda').manual_seed( random.randint(0, 9999) )
                    ).images[0]
        
        #image.save("astronaut_rides_horse.png")
        #image.save('output/image_'+str(i)+'_'+str(root_seed+seed).zfill(4)+'.png')
        image.save('images-IA-py/image_'+utilisateur+'_'+batch_count+'.png')
        return 1


