Update app.py

app.py

@@ -1,9 +1,33 @@
 import spaces
 import gradio as gr
-import transformers
-from transformers import AutoTokenizer
-from transformers import pipeline
+#import transformers
+#from transformers import AutoTokenizer
+#from transformers import pipeline
+from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
 
+base = "stabilityai/stable-diffusion-xl-base-1.0"
+repo = "ByteDance/SDXL-Lightning"
+ckpt = "sdxl_lightning_4step_unet.safetensors" # Use the correct ckpt for your step setting!
+
+# Load model.
+unet = UNet2DConditionModel.from_config(base, subfolder="unet").to("cuda", torch.float16)
+unet.load_state_dict(load_file(hf_hub_download(repo, ckpt), device="cuda"))
+pipe = StableDiffusionXLPipeline.from_pretrained(base, unet=unet, torch_dtype=torch.float16, variant="fp16").to("cuda")
+
+# Ensure sampler uses "trailing" timesteps.
+pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
+
+@spaces.GPU
+def run():
+    # Ensure using the same inference steps as the loaded model and CFG set to 0.
+    return pipe("A cat", num_inference_steps=4, guidance_scale=0).images[0].save("output.png")
+
+
+
+
+'''
 tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
 
 model = transformers.AutoModelForCausalLM.from_pretrained(
@@ -31,6 +55,9 @@ PROMPT_FOR_GENERATION_FORMAT = """{intro}
 example = "James decides to run 3 sprints 3 times a week. He runs 60 meters each sprint. How many total meters does he run a week? Explain before answering."
 fmt_ex = PROMPT_FOR_GENERATION_FORMAT.format(instruction=example)
 
+
+
+
 @spaces.GPU
 def run():
     with torch.autocast('cuda', dtype=torch.bfloat16):
@@ -39,12 +66,13 @@ def run():
                 max_new_tokens=100,
                 do_sample=True,
                 use_cache=True))
-
+'''
 
 
 with gr.Blocks() as app:
     btn = gr.Button()
-    outp=gr.Textbox()
+    #outp=gr.Textbox()
+    outp=gr.Image()
     btn.click(run,None,outp)
 app.launch()