Upload 2 files

app.py

@@ -6,7 +6,7 @@ import warnings
 warnings.filterwarnings('ignore')
 
 # Constants
-MODEL_ID = "CompVis/stable-diffusion-v1-4"  # Using SD 1.4 for better compatibility
+MODEL_ID = "runwayml/stable-diffusion-v1-5"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 DTYPE = torch.float32
 
@@ -16,17 +16,19 @@ class FooocusGenerator:
         self.load_models()
 
     def load_models(self):
-        self.pipe = StableDiffusionPipeline.from_pretrained(
-            MODEL_ID,
-            torch_dtype=DTYPE,
-            safety_checker=None,
-            requires_safety_checker=False
-        )
-        
-        if DEVICE == "cuda":
-            self.pipe.enable_attention_slicing()
-            self.pipe.enable_vae_slicing()
-        self.pipe = self.pipe.to(DEVICE)
+        try:
+            self.pipe = StableDiffusionPipeline.from_pretrained(
+                MODEL_ID,
+                torch_dtype=DTYPE,
+                safety_checker=None
+            )
+            
+            if DEVICE == "cuda":
+                self.pipe.enable_attention_slicing()
+            self.pipe = self.pipe.to(DEVICE)
+        except Exception as e:
+            print(f"Error loading model: {str(e)}")
+            raise e
 
     def generate_image(
         self,
@@ -65,15 +67,16 @@ class FooocusGenerator:
                 "Landscape": (768, 512)
             }.get(aspect_ratios_selection, (512, 512))
             
-            image = self.pipe(
-                prompt=processed_prompt,
-                negative_prompt=negative_prompt,
-                num_inference_steps=steps,
-                guidance_scale=guidance_scale,
-                generator=generator,
-                width=dimensions[0],
-                height=dimensions[1]
-            ).images[0]
+            with torch.inference_mode():
+                image = self.pipe(
+                    prompt=processed_prompt,
+                    negative_prompt=negative_prompt,
+                    num_inference_steps=steps,
+                    guidance_scale=guidance_scale,
+                    generator=generator,
+                    width=dimensions[0],
+                    height=dimensions[1]
+                ).images[0]
             
             return image
             

requirements.txt

@@ -1,14 +1,14 @@
-torch==2.0.1
-torchvision==0.15.2
-diffusers==0.11.1
-transformers==4.26.1
-accelerate==0.17.1
+torch==1.13.1
+torchvision==0.14.1
+diffusers==0.7.2
+transformers==4.25.1
+accelerate==0.15.0
 safetensors==0.3.1
-gradio==3.23.0
+gradio==3.16.2
 opencv-python==4.8.0.74
 einops==0.6.1
-pytorch-lightning==2.0.2
+pytorch-lightning==1.9.0
 omegaconf==2.3.0
-huggingface-hub==0.13.4
+huggingface-hub==0.11.1
 xformers==0.0.20
 triton==2.0.0