updates

app.py

@@ -1,6 +1,5 @@
 run_api = False
 SSD_1B = False
-
 import os
 
 # Use GPU
@@ -13,6 +12,7 @@ else:
     is_gpu = True
 print(is_gpu)
 
+
 from IPython.display import clear_output
 
 
@@ -38,6 +38,7 @@ def check_enviroment():
 # Call the function to check and install Packages if necessary
 check_enviroment()
 
+
 from IPython.display import clear_output
 import os
 import gradio as gr
@@ -70,51 +71,50 @@ else:
     # Uncomment the following line if you want to use CPU instead of GPU
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-if torch.cuda.is_available():
-    # Get the current directory
-    current_dir = os.getcwd()
-    model_path = os.path.join(current_dir)
 
-    # Set the cache path
-    cache_path = os.path.join(current_dir, "cache")
+# Get the current directory
+current_dir = os.getcwd()
+model_path = os.path.join(current_dir)
+
+# Set the cache path
+cache_path = os.path.join(current_dir, "cache")
 
-    if not SSD_1B:
+if not SSD_1B:
 
-        unet = UNet2DConditionModel.from_pretrained(
-            "latent-consistency/lcm-sdxl",
-            torch_dtype=torch.float16,
-            variant="fp16",
-            cache_dir=cache_path,
-        )
-        pipe = DiffusionPipeline.from_pretrained(
-            "stabilityai/stable-diffusion-xl-base-1.0",
-            unet=unet,
-            torch_dtype=torch.float16,
-            variant="fp16",
-            cache_dir=cache_path,
-        )
+    unet = UNet2DConditionModel.from_pretrained(
+        "latent-consistency/lcm-sdxl",
+        torch_dtype=torch.float16,
+        variant="fp16",
+        cache_dir=cache_path,
+    )
+    pipe = DiffusionPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-xl-base-1.0",
+        unet=unet,
+        torch_dtype=torch.float16,
+        variant="fp16",
+        cache_dir=cache_path,
+    )
 
-        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+    if torch.cuda.is_available():
         pipe.to("cuda")
-    else:
-        # SSD-1B
-        from diffusers import LCMScheduler, AutoPipelineForText2Image
-
-        pipe = AutoPipelineForText2Image.from_pretrained(
-            "segmind/SSD-1B",
-            torch_dtype=torch.float16,
-            variant="fp16",
-            cache_dir=cache_path,
-        )
-        pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+else:
+    # SSD-1B
+    from diffusers import LCMScheduler, AutoPipelineForText2Image
+
+    pipe = AutoPipelineForText2Image.from_pretrained(
+        "segmind/SSD-1B",
+        torch_dtype=torch.float16,
+        variant="fp16",
+        cache_dir=cache_path,
+    )
+    pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
+    if torch.cuda.is_available():
         pipe.to("cuda")
 
-        # load and fuse
-        pipe.load_lora_weights("latent-consistency/lcm-lora-ssd-1b")
-        pipe.fuse_lora()
-
-else:
-    pipe = None
+    # load and fuse
+    pipe.load_lora_weights("latent-consistency/lcm-lora-ssd-1b")
+    pipe.fuse_lora()
 
 
 def generate(
@@ -149,6 +149,25 @@ def generate(
 
 clear_output()
 
+from IPython.display import display
+
+
+def generate_image(prompt="A beautiful and sexy girl"):
+    # Generate the image using the prompt
+    generated_image = generate(
+        prompt=prompt,
+        negative_prompt="",
+        seed=0,
+        width=1024,
+        height=1024,
+        guidance_scale=0.0,
+        num_inference_steps=4,
+        secret_token="default_secret",  # Replace with your secret token
+    )
+    # Display the image in the Jupyter Notebook
+    display(generated_image)
+
+
 if not run_api:
     secret_token = gr.Text(
         label="Secret Token",

requirements.txt

@@ -5,3 +5,4 @@ invisible-watermark==0.2.0
 Pillow==10.1.0
 torch==2.1.0
 transformers==4.35.0
+ipython