Update app.py

app.py

@@ -4,61 +4,32 @@ import uuid
 import gradio as gr
 import numpy as np
 from PIL import Image
-import spaces
 import torch
-from diffusers import StableDiffusionPipeline, EulerAncestralDiscreteScheduler
+from huggingface_hub import login
+from diffusers import StableDiffusionPipeline
 from typing import Tuple
 
-def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
-    styles = {
-        "3840 x 2160": (
-            "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
-            "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
-        ),
-        "Style Zero": ("{prompt}", ""),
-    }
-    DEFAULT_STYLE_NAME = "3840 x 2160"
-
-    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
-    return p.replace("{prompt}", positive), n + negative
-
-def load_and_prepare_model(model_choice: str):
-    if model_choice == "DeepFloyd/IF-I-XL-v1.0":
-        model_id = "DeepFloyd/IF-I-XL-v1.0"
-        pipe = StableDiffusionPipeline.from_pretrained(
-            model_id,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            use_safetensors=True,
-            add_watermarker=False,
-        ).to(torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
-        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
-    else:
-        # Default to RealVisXL_V5.0_Lightning if no choice is made or another option is selected
-        model_id = "SG161222/RealVisXL_V5.0_Lightning"
-        pipe = StableDiffusionPipeline.from_pretrained(
-            model_id,
-            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            use_safetensors=True,
-            add_watermarker=False,
-        ).to(torch.device("cuda:0" if torch.cuda.is_available() else "cpu"))
-        pipe.scheduler = EulerAncestralDiscreteScheduler.from_config(pipe.scheduler.config)
+# Use the Hugging Face token from environment variable
+token = os.getenv("HF_TOKEN")  # Get token from environment variable
 
-    return pipe
+# Log in to Hugging Face using the token
+login(token=token)
 
-def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
-    if randomize_seed:
-        seed = random.randint(0, np.iinfo(np.int32).max)
-    return seed
+# Function to load and prepare the model
+def load_and_prepare_model(model_id: str):
+    # Load the model using the Stable Diffusion pipeline
+    pipe = StableDiffusionPipeline.from_pretrained(
+        model_id,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        use_safetensors=True
+    ).to(torch.device("cuda" if torch.cuda.is_available() else "cpu"))
 
-def save_image(img):
-    unique_name = str(uuid.uuid4()) + ".png"
-    img.save(unique_name)
-    return unique_name
+    return pipe
 
-@spaces.GPU(duration=60, enable_queue=True)
+# Function to generate the image
 def generate(
         prompt: str,
-        model_choice: str = "DeepFloyd/IF-I-XL-v1.0",  # Default model choice
+        model_choice: str,
         seed: int = 1,
         width: int = 1024,
         height: int = 1024,
@@ -66,9 +37,28 @@ def generate(
         num_inference_steps: int = 25,
         randomize_seed: bool = False,
 ):
-    model = load_and_prepare_model(model_choice)  # Load the selected model
-    seed = int(randomize_seed_fn(seed, randomize_seed))
-    generator = torch.Generator(device=model.device).manual_seed(seed)
+    # Set the seed
+    if randomize_seed:
+        seed = random.randint(0, np.iinfo(np.int32).max)
+
+    generator = torch.Generator(device="cuda" if torch.cuda.is_available() else "cpu").manual_seed(seed)
+
+    # Choose model based on the user's selection
+    model_id = "DeepFloyd/IF-I-XL-v1.0" if model_choice == "DeepFloyd/IF-I-XL-v1.0" else "RealVisXL_V5.0_Lightning"
+    model = load_and_prepare_model(model_id)
+
+    # Style functions (if needed)
+    def apply_style(style_name: str, positive: str, negative: str = "") -> Tuple[str, str]:
+        styles = {
+            "3840 x 2160": (
+                "hyper-realistic 8K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
+                "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
+            ),
+            "Style Zero": ("{prompt}", ""),
+        }
+        DEFAULT_STYLE_NAME = "3840 x 2160"
+        p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+        return p.replace("{prompt}", positive), n + negative
 
     positive_prompt, negative_prompt = apply_style("3840 x 2160", prompt)
 
@@ -83,36 +73,34 @@ def generate(
         "output_type": "pil",
     }
 
+    # Generate the image
     images = model(**options).images
-    image_paths = [save_image(img) for img in images]
-    return image_paths, seed
+    return images[0]
 
+# Gradio interface setup
 with gr.Blocks(theme="soft") as demo:
     # Centered text "SNAPSCRIBE" at the top of the screen
     gr.Markdown("<h1 style='text-align:center; color:white; font-weight:bold; text-decoration:underline;'>SNAPSCRIBE</h1>")
 
     with gr.Row():
         with gr.Column(scale=3):
-            # Dropdown for model selection
-            model_dropdown = gr.Dropdown(
-                label="Select Model",
-                choices=["DeepFloyd/IF-I-XL-v1.0", "SG161222/RealVisXL_V5.0_Lightning"],
-                value="DeepFloyd/IF-I-XL-v1.0",  # Default model choice
-            )
-            # Input Prompt text field
             prompt = gr.Textbox(
                 label="Input Prompt",
                 placeholder="Describe the image you want to create",
                 lines=2,
             )
+            # Drop down to select model
+            model_choice = gr.Dropdown(
+                choices=["DeepFloyd/IF-I-XL-v1.0", "RealVisXL_V5.0_Lightning"],
+                label="Choose a Model"
+            )
             run_button = gr.Button("Generate Image")
-            gr.Markdown("Developed using the DeepFloyd/IF-I-XL-v1.0 or RealVisXL_V5.0_Lightning model.", elem_id="model_info")
         with gr.Column(scale=7):
-            result = gr.Image(label="Generated Image", type="filepath")
+            result = gr.Image(label="Generated Image", type="pil")
 
     run_button.click(
         fn=generate,
-        inputs=[prompt, model_dropdown],
+        inputs=[prompt, model_choice],
         outputs=[result],
     )