Update app.py

app.py

@@ -12,7 +12,7 @@ from typing import Tuple
 
 bad_words = json.loads(os.getenv('BAD_WORDS', "[]"))
 bad_words_negative = json.loads(os.getenv('BAD_WORDS_NEGATIVE', "[]"))
-default_negative = os.getenv("default_negative","")
+default_negative = os.getenv("default_negative", "")
 
 def check_text(prompt, negative=""):
     for i in bad_words:
@@ -23,8 +23,7 @@ def check_text(prompt, negative=""):
             return True
     return False
 
-#Quality/Style-----------------------------------------------------------------------------------------------------------------------------------------------------------Quality/Style
-
+# Quality/Style
 style_list = [
     {
         "name": "3840 x 2160",
@@ -36,23 +35,19 @@ style_list = [
         "prompt": "hyper-realistic 4K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
         "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
     },
-
     {
         "name": "HD+",
         "prompt": "hyper-realistic 2K image of {prompt}. ultra-detailed, lifelike, high-resolution, sharp, vibrant colors, photorealistic",
         "negative_prompt": "cartoonish, low resolution, blurry, simplistic, abstract, deformed, ugly",
     },
-    
     {
         "name": "Style Zero",
         "prompt": "{prompt}",
         "negative_prompt": "",
     },
-    
 ]
 
-#Clgstyle--------------------------------------------------------------------------------------------------------------------------------------------------------------Clgstyle
-
+# Collage styles
 collage_style_list = [
     {
         "name": "Hi-Res",
@@ -109,9 +104,6 @@ collage_style_list = [
         "prompt": "thematic collage of {prompt}. cohesive theme, well-organized, matching colors, creative layout",
         "negative_prompt": "random, messy, unorganized, clashing colors",
     },
-
-#DuoTones by Canva --------------------------------------------------------------------------------------------------------------- Alters only the i++ Part / not Zero Tones
-    
     {
         "name": "Cherry",
         "prompt": "Duotone style Cherry tone applied to {prompt}",
@@ -182,11 +174,9 @@ collage_style_list = [
         "prompt": "{prompt}",
         "negative_prompt": "",
     },
-    
 ]
 
-#filters------------------------------------------------------------------------------------------------------------------------------------------------filters
-
+# Filters
 filters = {
     "Vivid": {
         "prompt": "extra vivid {prompt}",
@@ -212,13 +202,10 @@ filters = {
         "prompt": "{prompt} with a minimalist design",
         "negative_prompt": "cluttered, ornate"
     },
-
     "Zero filter": {
         "prompt": "{prompt}",
         "negative_prompt": ""
     },
-    
-    
 }
 
 styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
@@ -256,28 +243,42 @@ USE_TORCH_COMPILE = os.getenv("USE_TORCH_COMPILE", "0") == "1"
 ENABLE_CPU_OFFLOAD = os.getenv("ENABLE_CPU_OFFLOAD", "0") == "1"
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-#Compile
+# Set dtype based on device: half precision for CUDA, full precision for CPU
+dtype = torch.float16 if device.type == "cuda" else torch.float32
+
+# Load the pipeline
 if torch.cuda.is_available():
     pipe = StableDiffusionXLPipeline.from_pretrained(
-        "SG161222/RealVisXL_V5.0_Lightning", # / SG161222/RealVisXL_V4.0 / SG161222/RealVisXL_V4.0_Lightning 
-        torch_dtype=torch.float16,
+        "SG161222/RealVisXL_V5.0_Lightning",
+        torch_dtype=dtype,
         use_safetensors=True,
         add_watermarker=False
     ).to(device)
-
+    # Ensure the text encoder is in half precision to match the rest of the model
+    pipe.text_encoder = pipe.text_encoder.half()
+    
     if ENABLE_CPU_OFFLOAD:
         pipe.enable_model_cpu_offload()
     else:
         pipe.to(device)
         print("Loaded on Device!")
-
+    
     if USE_TORCH_COMPILE:
         pipe.unet = torch.compile(pipe.unet, mode="reduce-overhead", fullgraph=True)
         print("Model Compiled!")
+else:
+    # On CPU, load with float32
+    pipe = StableDiffusionXLPipeline.from_pretrained(
+        "SG161222/RealVisXL_V5.0_Lightning",
+        torch_dtype=dtype,
+        use_safetensors=True,
+        add_watermarker=False
+    ).to(device)
+    print("Running on CPU; model loaded in float32.")
 
 def save_image(img, path):
     img.save(path)
-#seeding
+
 def randomize_seed_fn(seed: int, randomize_seed: bool) -> int:
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
@@ -311,10 +312,10 @@ def generate(
         prompt, negative_prompt = apply_style(style, prompt, negative_prompt)
     
     seed = int(randomize_seed_fn(seed, randomize_seed))
-    generator = torch.Generator().manual_seed(seed)
+    generator = torch.Generator(device=device).manual_seed(seed)
 
     if not use_negative_prompt:
-        negative_prompt = ""  # type: ignore
+        negative_prompt = ""
     negative_prompt += default_negative    
 
     grid_sizes = {
@@ -342,7 +343,7 @@ def generate(
         "output_type": "pil",
     }
     
-    torch.cuda.empty_cache()  # Clear GPU memory
+    torch.cuda.empty_cache()  # Clear GPU memory if available
     images = pipe(**options).images
 
     grid_img = Image.new('RGB', (width * grid_size_x, height * grid_size_y))
@@ -355,18 +356,15 @@ def generate(
     return [unique_name], seed
 
 examples = [
-
-    "Chocolate dripping from a donut against a yellow background, in the style of brocore, hyper-realistic oil --ar 2:3 --q 2 --s 750 --v 5  --ar 2:3 --q 2 --s 750 --v 5",
+    "Chocolate dripping from a donut against a yellow background, in the style of brocore, hyper-realistic oil --ar 2:3 --q 2 --s 750 --v 5",
     "3d image, cute girl, in the style of Pixar --ar 1:2 --stylize 750, 4K resolution highlights, Sharp focus, octane render, ray tracing, Ultra-High-Definition, 8k, UHD, HDR, (Masterpiece:1.5), (best quality:1.5)",
     "Cold coffee in a cup bokeh --ar 85:128 --v 6.0 --style raw5,4k",
     "Food photography of a milk shake with flying strawberrys against a pink background, professionally studio shot with cinematic lighting. The image is in the style of a professional studio shot --ar 85:128 --v 6.0 --style raw"
-    
 ]
 
 css = '''
 .gradio-container{max-width: 888px !important}
 h1{text-align:center}
-
 .submit-btn {
     background: linear-gradient(90deg, #4B79A1 0%, #283E51 100%) !important;
     border: none !important;