Create app.py

app.py

@@ -0,0 +1,30 @@
+import gradio as gr
+from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel
+import torch
+
+# load pipeline
+model_id = "stabilityai/stable-diffusion-xl-base-1.0"
+pipe = StableDiffusionXLPipeline.from_pretrained(model_id, torch_dtype=torch.float16, variant="fp16", use_safetensors=True).to("cuda")
+
+# load finetuned model
+unet_id = "mhdang/dpo-sdxl-text2image-v1"
+unet = UNet2DConditionModel.from_pretrained(unet_id, subfolder="unet", torch_dtype=torch.float16)
+pipe.unet = unet
+pipe = pipe.to("cuda")
+
+pipe.enable_model_cpu_offload() 
+pipe.enable_vae_slicing()
+
+def infer(prompt):
+    image = pipe(prompt, guidance_scale=7.5).images[0].resize((512,512))
+    return image
+
+gr.Interface(
+    fn = infer,
+    inputs = [
+        gr.Textbox(value="Two cats playing chess on a tree branch")
+    ],
+    outputs = [
+        gr.Image()
+    ]
+).launch()