Create model.py

model.py

@@ -0,0 +1,18 @@
+from diffusers import DiffusionPipeline
+import torch
+
+# Replace with your model path if it's not from the Hugging Face model hub
+# model_id = "https://api-inference.huggingface.co/models/awabxnero/imagen"  # Your Hugging Face model repo
+pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell")
+pipe = pipe.to("cuda")  # Move model to GPU if available
+
+# Function to generate image from prompt
+def generate_image(prompt):
+    # Generate image
+    image = pipe(prompt).images[0]
+    image.show()  # Show the image (you can also save or return the image)
+    return image
+
+if __name__ == "__main__":
+    prompt = "A beautiful sunset over the ocean"
+    generate_image(prompt)