add login

app.py

@@ -4,12 +4,16 @@ from diffusers import DiffusionPipeline, StableDiffusionPipeline, StableDiffusio
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from pathlib import Path
 from safetensors.torch import load_file
-from huggingface_hub import hf_hub_download, list_models
+from huggingface_hub import hf_hub_download
 from PIL import Image
 import matplotlib.pyplot as plt
 from matplotlib.colors import hex2color
 import stone
 import os
+import spaces
+
+from huggingface_hub import login
+login()
 
 # Define model initialization functions
 def load_model(model_name, token=None):
@@ -59,10 +63,10 @@ def load_model(model_name, token=None):
 default_model = "stabilityai/sdxl-turbo"
 pipeline_text2image = load_model(default_model)
 
-@gr.outputs.Image(type="pil")
-def getimgen(prompt, model_name, token):
+@spaces.GPU
+def getimgen(prompt, model_name):
     global pipeline_text2image
-    pipeline_text2image = load_model(model_name, token)
+    pipeline_text2image = load_model(model_name)
     if model_name == "stabilityai/sdxl-turbo":
         return pipeline_text2image(prompt=prompt, guidance_scale=0.0, num_inference_steps=2).images[0]
     elif model_name == "runwayml/stable-diffusion-v1-5":
@@ -78,7 +82,7 @@ def getimgen(prompt, model_name, token):
 blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
 blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large", torch_dtype=torch.float16).to("cuda")
 
-@gr.outputs.Textbox(type="str")
+@spaces.GPU
 def blip_caption_image(image, prefix):
     inputs = blip_processor(image, prefix, return_tensors="pt").to("cuda", torch.float16)
     out = blip_model.generate(**inputs)
@@ -117,11 +121,13 @@ def skintoneplot(hex_codes):
             ax.add_patch(plt.Rectangle((0, 0), 1, 1, color=sorted_hex_codes[i]))
     return fig
 
-@gr.outputs.Image(type="pil")
-def generate_images_plots(prompt, model_name, token):
+@spaces.GPU
+def generate_images_plots(prompt, model_name):
+    global pipeline_text2image
+    pipeline_text2image = load_model(model_name)
     foldername = "temp"
     Path(foldername).mkdir(parents=True, exist_ok=True)
-    images = [getimgen(prompt, model_name, token) for _ in range(10)]
+    images = [getimgen(prompt, model_name) for _ in range(10)]
     genders = []
     skintones = []
     for image, i in zip(images, range(10)):
@@ -139,9 +145,6 @@ def generate_images_plots(prompt, model_name, token):
 
 with gr.Blocks(title="Skin Tone and Gender bias in Text to Image Models") as demo:
     gr.Markdown("# Skin Tone and Gender bias in Text to Image Models")
-    gr.LoginButton()  # Add a login button for Hugging Face
-    profile = gr.State()
-    token = gr.State()
     model_dropdown = gr.Dropdown(
         label="Choose a model", 
         choices=[
@@ -163,10 +166,12 @@ with gr.Blocks(title="Skin Tone and Gender bias in Text to Image Models") as dem
         object_fit="contain", 
         height="auto"
     )
+    gr.LoginButton() 
+    gr.Markdown('### You need to log in to your Hugging Face account to run Stable Diffusion 3')
     btn = gr.Button("Generate images", scale=0)
     with gr.Row(equal_height=True):
         skinplot = gr.Plot(label="Skin Tone")
         genplot = gr.Plot(label="Gender")
-    btn.click(generate_images_plots, inputs=[prompt, model_dropdown, token], outputs=[gallery, skinplot, genplot])
+    btn.click(generate_images_plots, inputs=[prompt, model_dropdown], outputs=[gallery, skinplot, genplot])
 
-demo.launch()
+demo.launch(debug=True)