import gradio as gr
from diffusers import DiffusionPipeline
import torch

def get_device():
  if torch.cuda.is_available():
    return "cuda"
  else:
    return "cpu"

def generate_image(prompt):
  pipe_id = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
  pipe = DiffusionPipeline.from_pretrained(pipe_id, torch_dtype=torch.float16).to("cuda")
  pipe.load_lora_weights("timdpaep/t1m")
  prompt = "professional photo, closeup photo of t1mLora, wearing black sweater, nature, gloomy, cloudy weather, bokeh <lora:t1m01:1>"

  lora_scale= 0.9
  image = pipe(
    prompt, num_inference_steps=10, cross_attention_kwargs={"scale": lora_scale}, generator=torch.manual_seed(0)
  ).to(get_device()).images[0]
  return image


iface = gr.Interface(fn=generate_image, inputs="textbox", outputs="image")
iface.launch()