Update app.py

app.py

@@ -2,32 +2,40 @@ import gradio as gr
 
 import subprocess
 
-subprocess.check_call(["pip", "install", "transformers"])
+subprocess.check_call(["pip", "install", "safetensors"])
 subprocess.check_call(["pip", "install", "torch"])
+subprocess.check_call(["pip", "install", "diffusers"])
 
+import torch
+from diffusers import StableDiffusionXLPipeline, UNet2DConditionModel, EulerDiscreteScheduler
+from huggingface_hub import hf_hub_download
+from safetensors.torch import load_file
 
-from transformers import pipeline
+base = "stabilityai/stable-diffusion-xl-base-1.0"
+repo = "ByteDance/SDXL-Lightning"
+ckpt = "sdxl_lightning_2step_unet.safetensors" # Use the correct ckpt for your step setting!
 
-summariser = pipeline("text2text-generation", model="balaramas/mbart-enhiriser")
+# 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")
+pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
 
+def generate_image(text):
 
-def summarise(Input):
-        
-    res = summariser(Input, max_length=120)[0]['generated_text']
-    return res
-
-description = """
-<p>
-<center>
-Multi-domain Summarisation Between English and Hindi
-</center>
-</p>
-"""
+    image = pipe("krishna", num_inference_steps=2, guidance_scale=0).images[0]
+    
+    return image
 
+# Create a Gradio interface
 iface = gr.Interface(
-    fn=summarise,
-    inputs="text",
-    title="🌸English to Hindi Summariser🌸",
-    outputs="text")
-
+    fn=generate_image,
+    inputs=gr.inputs.Textbox(lines=5, label="Enter a description for the image"),
+    outputs="image",
+    title="Text to Image Generation",
+    description="Enter a text description and get an image.",
+    theme="compact"
+)
+
+# Launch the Gradio app
 iface.launch()