add login

app.py

@@ -4,19 +4,15 @@ 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
+from huggingface_hub import hf_hub_download, list_models
 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):
+def load_model(model_name, token=None):
     if model_name == "stabilityai/sdxl-turbo":
         pipeline = DiffusionPipeline.from_pretrained(
             model_name, 
@@ -48,9 +44,12 @@ def load_model(model_name):
             variant="fp16"
         ).to("cuda")
     elif model_name == "stabilityai/stable-diffusion-3-medium-diffusers":
+        if token is None:
+            raise ValueError("Hugging Face token is required to access this model")
         pipeline = StableDiffusion3Pipeline.from_pretrained(
             model_name, 
-            torch_dtype=torch.float16
+            torch_dtype=torch.float16,
+            use_auth_token=token
         ).to("cuda")
     else:
         raise ValueError("Unknown model name")
@@ -60,10 +59,10 @@ def load_model(model_name):
 default_model = "stabilityai/sdxl-turbo"
 pipeline_text2image = load_model(default_model)
 
-@spaces.GPU
-def getimgen(prompt, model_name):
+@gr.outputs.Image(type="pil")
+def getimgen(prompt, model_name, token):
     global pipeline_text2image
-    pipeline_text2image = load_model(model_name)
+    pipeline_text2image = load_model(model_name, token)
     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":
@@ -79,7 +78,7 @@ def getimgen(prompt, model_name):
 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")
 
-@spaces.GPU
+@gr.outputs.Textbox(type="str")
 def blip_caption_image(image, prefix):
     inputs = blip_processor(image, prefix, return_tensors="pt").to("cuda", torch.float16)
     out = blip_model.generate(**inputs)
@@ -118,13 +117,11 @@ def skintoneplot(hex_codes):
             ax.add_patch(plt.Rectangle((0, 0), 1, 1, color=sorted_hex_codes[i]))
     return fig
 
-@spaces.GPU
-def generate_images_plots(prompt, model_name):
-    global pipeline_text2image
-    pipeline_text2image = load_model(model_name)
+@gr.outputs.Image(type="pil")
+def generate_images_plots(prompt, model_name, token):
     foldername = "temp"
     Path(foldername).mkdir(parents=True, exist_ok=True)
-    images = [getimgen(prompt, model_name) for _ in range(10)]
+    images = [getimgen(prompt, model_name, token) for _ in range(10)]
     genders = []
     skintones = []
     for image, i in zip(images, range(10)):
@@ -142,6 +139,9 @@ def generate_images_plots(prompt, model_name):
 
 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=[
@@ -167,6 +167,6 @@ with gr.Blocks(title="Skin Tone and Gender bias in Text to Image Models") as dem
     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], outputs=[gallery, skinplot, genplot])
+    btn.click(generate_images_plots, inputs=[prompt, model_dropdown, token], outputs=[gallery, skinplot, genplot])
 
-demo.launch(debug=True)
+demo.launch()