Update app.py

app.py

@@ -1,24 +1,34 @@
 from fastai.vision.all import *
 import gradio as gr
+from huggingface_hub import login
+login()
+import torch
+from diffusers import FluxPipeline
+
+pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-dev", torch_dtype=torch.bfloat16)
+pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
+
+prompt = "A cat holding a sign that says hello world"
+image = pipe(
+    prompt,
+    height=1024,
+    width=1024,
+    guidance_scale=3.5,
+    num_inference_steps=50,
+    max_sequence_length=512,
+    generator=torch.Generator("cpu").manual_seed(0)
+).images[0]
+image.save("flux-dev.png")
 
-im1 = PILImage.create('Brittlebush-Encelia-Farinosa-desert-horizon-nursery.jpg')
-im2 = PILImage.create('download.webp')
 
 learn = load_learner('export.pkl')
 
-pred_class,pred_idx,probabilities = learn.predict(im1)
-pred_class, pred_idx, probabilities
-
-pred_class,pred_idx,probabilities = learn.predict(im2)
-pred_class, pred_idx, probabilities
-
 categories = ('balsamroot', 'bladderpod', 'blazing star', 'bristlecone pine flowers', 'brittlebrush')
 def classify_image(img):
     pred, idx, probs = learn.predict(img)
     return dict(zip(categories, map(float, probs)))
 
 
-classify_image(im1), classify_image(im2)
 
 image=gr.Image(height = 192, width = 192)
 label = gr.Label()