Rename app.py to app_cpu.py

app.py → app_cpu.py

@@ -1,12 +1,16 @@
 import gradio as gr
 import torch
-from transformers import CLIPTextModel, CLIPTokenizer
-from diffusers import StableDiffusionPipeline
+from diffusers import DiffusionPipeline
+from PIL import Image
 
-# Load the model and tokenizer
+# Load the model
 model_id = "stabilityai/stable-diffusion-xl-base-1.0"
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
-pipe = pipe.to("cpu")
+pipe = DiffusionPipeline.from_pretrained(
+    model_id,
+    torch_dtype=torch.float32,  # Use float32 for CPU
+    use_safetensors=True
+)
+pipe.to("cpu")
 
 def generate_image(prompt, negative_prompt, size):
     if not prompt:
@@ -16,13 +20,16 @@ def generate_image(prompt, negative_prompt, size):
         
     width, height = map(int, size.split('x'))
     generator = torch.Generator("cpu").manual_seed(42)
-    
-    # Generate the image
-    result = pipe(prompt, height=height, width=width, negative_prompt=negative_prompt, generator=generator)
-    
-    if result is not None and 'images' in result:
-        return result.images[0]
-    else:
+
+    try:
+        result = pipe(prompt=prompt, height=height, width=width, negative_prompt=negative_prompt, generator=generator)
+        if result and hasattr(result, 'images') and len(result.images) > 0:
+            return result.images[0]
+        else:
+            print("Error: No images in the result or result is None")
+            return None
+    except Exception as e:
+        print(f"Error occurred: {e}")
         return None
 
 with gr.Blocks() as demo:
@@ -41,6 +48,7 @@ with gr.Blocks() as demo:
 
 demo.launch()
 
+
 # import gradio as gr
 # import torch
 # from transformers import CLIPTextModel, CLIPTokenizer