Upload app.py

app.py

@@ -390,42 +390,6 @@ DEVICE = None
 CONFIG = None
 
 
-def save_to_hub(save_name, repo_id=None):
-    """Save model checkpoint to HuggingFace Hub for persistence."""
-    global MODEL, TEXT_ENCODER, CONFIG
-    
-    if MODEL is None:
-        return "❌ No model loaded to save"
-    
-    try:
-        from huggingface_hub import HfApi, upload_file
-        import tempfile
-        
-        # Save to temp file
-        with tempfile.NamedTemporaryFile(suffix='.pt', delete=False) as f:
-            torch.save({
-                "model_state_dict": MODEL.state_dict(),
-                "text_encoder_state_dict": TEXT_ENCODER.state_dict(),
-                "config": CONFIG
-            }, f.name)
-            temp_path = f.name
-        
-        # Upload to Hub (same Space repo)
-        api = HfApi()
-        api.upload_file(
-            path_or_fileobj=temp_path,
-            path_in_repo=f"checkpoints/{save_name}.pt",
-            repo_id=repo_id or "Spanicin/candlestick-diffusion",
-            repo_type="space"
-        )
-        
-        os.unlink(temp_path)
-        return f"✅ Model saved to Hub: checkpoints/{save_name}.pt"
-    
-    except Exception as e:
-        return f"❌ Failed to save to Hub: {str(e)}"
-
-
 def load_model(checkpoint_path=None):
     global MODEL, TEXT_ENCODER, DIFFUSION, DEVICE, CONFIG
     
@@ -478,33 +442,25 @@ def generate_dataset_ui(num_samples, image_size):
         import os
         import json
         import random
-        from PIL import Image, ImageDraw
+        from PIL import Image
+        import matplotlib.pyplot as plt
+        from matplotlib.patches import Rectangle
+        import io
         
         output_dir = "./dataset"
         os.makedirs(output_dir, exist_ok=True)
         os.makedirs(os.path.join(output_dir, "images"), exist_ok=True)
         
+        bg_color = "#1a1a2e"
+        bullish_color = "#00ff88"
+        bearish_color = "#ff4466"
         num_candles = 20
-        labels = {}
-        
-        # Colors
-        bg_color = (26, 26, 46)
-        bullish_color = (0, 255, 136)
-        bearish_color = (255, 68, 102)
-        
-        patterns = ["bullish", "bearish", "sideways"]
-        volatilities = {"low": 1.0, "medium": 3.0, "high": 6.0}
         
-        for i in range(int(num_samples)):
-            # Generate candle data
-            pattern = random.choice(patterns)
-            vol_name = random.choice(list(volatilities.keys()))
-            vol = volatilities[vol_name]
-            
+        def generate_candles(pattern, vol):
             candles = []
             price = 100 if pattern != "bearish" else 150
             
-            for j in range(num_candles):
+            for i in range(num_candles):
                 if pattern == "bullish":
                     trend = random.uniform(0.5, 2.0)
                     o = price + random.gauss(0, vol)
@@ -513,7 +469,7 @@ def generate_dataset_ui(num_samples, image_size):
                     trend = random.uniform(0.5, 2.0)
                     o = price + random.gauss(0, vol)
                     c = o - random.uniform(0, vol*2) - trend
-                else:
+                else:  # sideways
                     o = price + random.gauss(0, vol)
                     c = o + random.gauss(0, vol)
                 
@@ -521,62 +477,63 @@ def generate_dataset_ui(num_samples, image_size):
                 l = min(o, c) - random.uniform(0, vol)
                 candles.append({"o": o, "h": h, "l": l, "c": c})
                 price = c
-            
-            # Render with PIL (no matplotlib)
-            img = Image.new('RGB', (image_size, image_size), bg_color)
-            draw = ImageDraw.Draw(img)
+            return candles
+        
+        def render(candles):
+            fig, ax = plt.subplots(figsize=(image_size/100, image_size/100), dpi=100)
+            fig.patch.set_facecolor(bg_color)
+            ax.set_facecolor(bg_color)
             
             highs = [c["h"] for c in candles]
             lows = [c["l"] for c in candles]
-            price_min, price_max = min(lows), max(highs)
-            price_range = price_max - price_min or 1
-            
-            padding = 10
-            chart_width = image_size - 2 * padding
-            chart_height = image_size - 2 * padding
-            candle_width = chart_width // (num_candles + 2)
+            price_min, price_max = min(lows)*0.98, max(highs)*1.02
             
-            for j, c in enumerate(candles):
-                x = padding + (j + 1) * candle_width
-                
-                # Scale prices to pixels
-                def scale_y(p):
-                    return padding + chart_height - int((p - price_min) / price_range * chart_height)
-                
-                y_open = scale_y(c["o"])
-                y_close = scale_y(c["c"])
-                y_high = scale_y(c["h"])
-                y_low = scale_y(c["l"])
-                
+            for i, c in enumerate(candles):
                 color = bullish_color if c["c"] >= c["o"] else bearish_color
-                
-                # Draw wick
-                draw.line([(x, y_high), (x, y_low)], fill=color, width=1)
-                
-                # Draw body
-                body_top = min(y_open, y_close)
-                body_bottom = max(y_open, y_close)
-                if body_bottom - body_top < 2:
-                    body_bottom = body_top + 2
-                draw.rectangle(
-                    [(x - candle_width//3, body_top), (x + candle_width//3, body_bottom)],
-                    fill=color
-                )
+                ax.plot([i, i], [c["l"], c["h"]], color=color, linewidth=1)
+                body_bottom = min(c["o"], c["c"])
+                body_height = abs(c["c"] - c["o"]) or 0.1
+                rect = Rectangle((i-0.3, body_bottom), 0.6, body_height, facecolor=color)
+                ax.add_patch(rect)
+            
+            ax.set_xlim(-1, len(candles))
+            ax.set_ylim(price_min, price_max)
+            ax.axis("off")
+            
+            buf = io.BytesIO()
+            plt.savefig(buf, format="png", facecolor=bg_color, bbox_inches="tight", pad_inches=0.1)
+            plt.close(fig)
+            buf.seek(0)
+            
+            img = Image.open(buf).convert("RGB")
+            return img.resize((image_size, image_size), Image.Resampling.LANCZOS)
+        
+        patterns = ["bullish", "bearish", "sideways"]
+        volatilities = {"low": 1.0, "medium": 3.0, "high": 6.0}
+        labels = {}
+        
+        for i in range(int(num_samples)):
+            pattern = random.choice(patterns)
+            vol_name = random.choice(list(volatilities.keys()))
+            vol = volatilities[vol_name]
+            
+            candles = generate_candles(pattern, vol)
+            img = render(candles)
             
-            # Save image
             filename = f"chart_{i:06d}.png"
             img.save(os.path.join(output_dir, "images", filename))
             labels[filename] = f"{pattern} trend {vol_name} volatility"
+            
+            if i % 500 == 0:
+                print(f"Generated {i}/{num_samples}")
         
-        # Save labels
         with open(os.path.join(output_dir, "labels.json"), "w") as f:
             json.dump(labels, f)
         
-        return f"✅ Generated {int(num_samples)} samples in ./dataset"
+        return f"✅ Generated {num_samples} samples in ./dataset"
     
     except Exception as e:
-        import traceback
-        return f"❌ Failed: {str(e)}\n{traceback.format_exc()}"
+        return f"❌ Failed: {str(e)}"
 
 
 # ============== Gradio Interface ==============
@@ -625,20 +582,12 @@ def train_model(data_path, epochs, batch_size, learning_rate, image_size, save_n
     global MODEL, TEXT_ENCODER, DIFFUSION, DEVICE, CONFIG
     
     try:
-        # Clear GPU memory
-        import gc
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
-        
         # Setup
         DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        
-        # Use smaller model for T4 GPU
         CONFIG = {
-            "base_channels": 48,  # Reduced from 64
-            "channel_mults": (1, 2, 4),  # Keep same
-            "context_dim": 192,  # Reduced from 256
+            "base_channels": 64,
+            "channel_mults": (1, 2, 4),
+            "context_dim": 256,
             "image_size": image_size,
             "timesteps": 1000
         }
@@ -674,14 +623,13 @@ def train_model(data_path, epochs, batch_size, learning_rate, image_size, save_n
         logs = [f"🚀 Training started on {DEVICE}"]
         logs.append(f"📊 Model parameters: {num_params:,}")
         logs.append(f"📁 Training samples: {len(train_dataset)}")
-        logs.append(f"🖼️ Image size: {image_size}x{image_size}")
-        logs.append(f"📦 Batch size: {batch_size}")
         logs.append("-" * 40)
         
-        for epoch in range(int(epochs)):
+        total_steps = epochs * len(train_loader)
+        current_step = 0
+        
+        for epoch in range(epochs):
             epoch_loss = 0
-            batch_count = 0
-            
             for images, texts in train_loader:
                 images = images.to(DEVICE)
                 context = TEXT_ENCODER(texts, DEVICE)
@@ -693,28 +641,10 @@ def train_model(data_path, epochs, batch_size, learning_rate, image_size, save_n
                 optimizer.step()
                 
                 epoch_loss += loss.item()
-                batch_count += 1
-                
-                # Clear cache periodically
-                if batch_count % 50 == 0:
-                    if torch.cuda.is_available():
-                        torch.cuda.empty_cache()
+                current_step += 1
             
-            avg_loss = epoch_loss / max(batch_count, 1)
-            logs.append(f"Epoch {epoch+1}/{int(epochs)}: loss = {avg_loss:.4f}")
-            
-            # Save checkpoint every 10 epochs
-            if (epoch + 1) % 10 == 0:
-                print(f"Epoch {epoch+1}: loss = {avg_loss:.4f}")
-                checkpoint_path = f"checkpoints/{save_name}_epoch{epoch+1}.pt"
-                os.makedirs("checkpoints", exist_ok=True)
-                torch.save({
-                    "model_state_dict": MODEL.state_dict(),
-                    "text_encoder_state_dict": TEXT_ENCODER.state_dict(),
-                    "config": CONFIG,
-                    "epoch": epoch + 1
-                }, checkpoint_path)
-                logs.append(f"💾 Checkpoint saved: {checkpoint_path}")
+            avg_loss = epoch_loss / len(train_loader)
+            logs.append(f"Epoch {epoch+1}/{epochs}: loss = {avg_loss:.4f}")
         
         # Save model
         MODEL.eval()
@@ -729,26 +659,10 @@ def train_model(data_path, epochs, batch_size, learning_rate, image_size, save_n
         logs.append("-" * 40)
         logs.append(f"✅ Model saved to {save_path}")
         
-        # Also try to save to Hub for persistence
-        try:
-            from huggingface_hub import HfApi
-            api = HfApi()
-            api.upload_file(
-                path_or_fileobj=save_path,
-                path_in_repo=f"checkpoints/{save_name}.pt",
-                repo_id="Spanicin/candlestick-diffusion",
-                repo_type="space"
-            )
-            logs.append(f"☁️ Model uploaded to Hub (persistent)")
-        except Exception as hub_error:
-            logs.append(f"⚠️ Could not upload to Hub: {hub_error}")
-            logs.append("   Model saved locally but may be lost on restart")
-        
         return "\n".join(logs)
     
     except Exception as e:
-        import traceback
-        return f"❌ Training failed: {str(e)}\n{traceback.format_exc()}"
+        return f"❌ Training failed: {str(e)}"
 
 
 def load_checkpoint(checkpoint_file):