Update app.py

app.py

@@ -4,32 +4,50 @@ import uuid
 import gradio as gr
 import numpy as np
 from PIL import Image
+import spaces
 import torch
-from huggingface_hub import login
-from diffusers import StableDiffusionPipeline
+from diffusers import StableDiffusionXLPipeline, EulerAncestralDiscreteScheduler
 from typing import Tuple
 
-# Use the Hugging Face token from environment variable
-token = os.getenv("HF_TOKEN")  # Get token from environment variable
+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"
 
-# Log in to Hugging Face using the token
-login(token=token)
+    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
+    return p.replace("{prompt}", positive), n + negative
 
-# 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(
+def load_and_prepare_model():
+    model_id = "SG161222/RealVisXL_V5.0_Lightning"
+    pipe = StableDiffusionXLPipeline.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"))
-
+        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)
     return pipe
 
-# Function to generate the image
+model = load_and_prepare_model()
+
+def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
+    if randomize_seed:
+        seed = random.randint(0, np.iinfo(np.int32).max)
+    return seed
+
+def save_image(img):
+    unique_name = str(uuid.uuid4()) + ".png"
+    img.save(unique_name)
+    return unique_name
+
+@spaces.GPU(duration=60, enable_queue=True)
 def generate(
         prompt: str,
-        model_choice: str,
         seed: int = 1,
         width: int = 1024,
         height: int = 1024,
@@ -37,28 +55,10 @@ def generate(
         num_inference_steps: int = 25,
         randomize_seed: bool = False,
 ):
-    # Set the seed
-    if randomize_seed:
-        seed = random.randint(0, np.iinfo(np.int32).max)
+    global model
 
-    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
+    seed = int(randomize_seed_fn(seed, randomize_seed))
+    generator = torch.Generator(device=model.device).manual_seed(seed)
 
     positive_prompt, negative_prompt = apply_style("3840 x 2160", prompt)
 
@@ -73,15 +73,14 @@ def generate(
         "output_type": "pil",
     }
 
-    # Generate the image
     images = model(**options).images
-    return images[0]
+    image_path = save_image(images[0])  # Saving the first generated image
+    return image_path
 
-# 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):
             prompt = gr.Textbox(
@@ -89,37 +88,36 @@ with gr.Blocks(theme="soft") as demo:
                 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 RealVisXL_V5.0_Lightning model.", elem_id="model_info")
         with gr.Column(scale=7):
-            result = gr.Image(label="Generated Image", type="pil")
+            result_image = gr.Image(label="Generated Image", type="filepath")
 
     run_button.click(
         fn=generate,
-        inputs=[prompt, model_choice],
-        outputs=[result],
+        inputs=[prompt],
+        outputs=[result_image],
     )
 
     # Footer with custom style and text
-    gr.Markdown("""
+    gr.HTML("""
     <style>
-    .footer {
-        position: relative;
-        left: 0;
-        bottom: 0;
-        width: 100%;
-        background-color: white;
-        color: black;
-        text-align: center;
-    }
+        .footer {
+            position: relative;
+            left: 0;
+            bottom: 0;
+            width: 100%;
+            background-color: white;
+            color: black;
+            text-align: center;
+            padding: 10px;
+            margin-top: 20px;
+        }
     </style>
     <div class="footer">
-    <p>Developed with ❤ by Aklavya(Bucky)</p>
+        <p>Developed with ❤ by Aklavya(Bucky)</p>
     </div>
     """)
 
 demo.launch()
+