Upload app.py

app.py

@@ -0,0 +1,45 @@
+import streamlit as st
+from diffusers import UNet2DConditionModel, DiffusionPipeline, LCMScheduler
+import torch
+from PIL import Image
+
+# Function to generate and display image
+def generate_and_display_image(prompt):
+    # Initialize the UNet model
+    unet = UNet2DConditionModel.from_pretrained("path/to/fine-tuned/weight", torch_dtype=torch.float16, variant="fp16")
+
+    # Initialize the diffusion pipeline
+    pipeline = DiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", unet=unet, torch_dtype=torch.float16, variant="fp16")
+    pipeline.safety_checker = None
+    pipeline.requires_safety_checker = False
+
+    # Set the loaded scheduler in the pipeline
+    pipeline.scheduler = LCMScheduler.from_config(pipeline.scheduler.config)
+    pipeline.to("cuda")
+
+    # Set the number of inference steps
+    inference_steps = 4
+
+    # Generate image
+    image = pipeline(prompt, num_inference_steps=inference_steps, guidance_scale=2).images[0]
+    image = image.resize((512, 512))
+
+    # Display the generated image
+    st.image(image, caption="Generated Image", use_column_width=True)
+
+# Main function
+def main():
+    st.title("Image Generation with Diffusion Models")
+
+    # Input prompt
+    prompt = st.text_input("Enter your prompt")
+
+    # Button to generate and display image
+    if st.button("Generate Image"):
+        if prompt:
+            generate_and_display_image(prompt)
+        else:
+            st.warning("Please provide a prompt.")
+
+if __name__ == "__main__":
+    main()